From a30c3ad61e717a6325c392ce661bd38fa7044556 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Tue, 2 Feb 2016 00:51:17 -0800
Subject: [PATCH] etcdserver: reject v3 txns with duplicate put keys

An API check to support PR #4363; bad requests didn't return an error.
---
 etcdserver/api/v3rpc/error.go |  1 +
 etcdserver/api/v3rpc/key.go   | 27 +++++++++++++++++++++++++++
 integration/v3_grpc_test.go   | 26 +++++++++++++++++++++++---
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/etcdserver/api/v3rpc/error.go b/etcdserver/api/v3rpc/error.go
index 3b266ac1f95c..ab01edf75912 100644
--- a/etcdserver/api/v3rpc/error.go
+++ b/etcdserver/api/v3rpc/error.go
@@ -23,6 +23,7 @@ import (
 var (
 	ErrEmptyKey      = grpc.Errorf(codes.InvalidArgument, "key is not provided")
 	ErrTooManyOps    = grpc.Errorf(codes.InvalidArgument, "too many operations in txn request")
+	ErrDuplicateKey  = grpc.Errorf(codes.InvalidArgument, "duplicate key given in txn request")
 	ErrCompacted     = grpc.Errorf(codes.OutOfRange, storage.ErrCompacted.Error())
 	ErrFutureRev     = grpc.Errorf(codes.OutOfRange, storage.ErrFutureRev.Error())
 	ErrLeaseNotFound = grpc.Errorf(codes.NotFound, "requested lease not found")
diff --git a/etcdserver/api/v3rpc/key.go b/etcdserver/api/v3rpc/key.go
index 52a8bb564a3c..349a01a95950 100644
--- a/etcdserver/api/v3rpc/key.go
+++ b/etcdserver/api/v3rpc/key.go
@@ -176,13 +176,40 @@ func checkTxnRequest(r *pb.TxnRequest) error {
 			return err
 		}
 	}
+	if err := checkRequestPutKeys(r.Success); err != nil {
+		return err
+	}
 
 	for _, u := range r.Failure {
 		if err := checkRequestUnion(u); err != nil {
 			return err
 		}
 	}
+	if err := checkRequestPutKeys(r.Failure); err != nil {
+		return err
+	}
+
+	return nil
+}
 
+// checkRequestPutKeys gives ErrDuplicateKey if the same key is put twice
+func checkRequestPutKeys(reqs []*pb.RequestUnion) error {
+	keys := make(map[string]struct{})
+	for _, requ := range reqs {
+		tv, ok := requ.Request.(*pb.RequestUnion_RequestPut)
+		if !ok {
+			continue
+		}
+		preq := tv.RequestPut
+		if preq == nil {
+			continue
+		}
+		key := string(preq.Key)
+		if _, ok := keys[key]; ok {
+			return ErrDuplicateKey
+		}
+		keys[key] = struct{}{}
+	}
 	return nil
 }
 
diff --git a/integration/v3_grpc_test.go b/integration/v3_grpc_test.go
index a898e24ac619..61388f30b506 100644
--- a/integration/v3_grpc_test.go
+++ b/integration/v3_grpc_test.go
@@ -84,12 +84,19 @@ func TestV3TxnTooManyOps(t *testing.T) {
 
 	kvc := clus.RandClient().KV
 
+	// unique keys
+	i := new(int)
+	keyf := func() []byte {
+		*i++
+		return []byte(fmt.Sprintf("key-%d", i))
+	}
+
 	addCompareOps := func(txn *pb.TxnRequest) {
 		txn.Compare = append(txn.Compare,
 			&pb.Compare{
 				Result: pb.Compare_GREATER,
 				Target: pb.Compare_CREATE,
-				Key:    []byte("bar"),
+				Key:    keyf(),
 			})
 	}
 	addSuccessOps := func(txn *pb.TxnRequest) {
@@ -97,7 +104,7 @@ func TestV3TxnTooManyOps(t *testing.T) {
 			&pb.RequestUnion{
 				Request: &pb.RequestUnion_RequestPut{
 					RequestPut: &pb.PutRequest{
-						Key:   []byte("bar"),
+						Key:   keyf(),
 						Value: []byte("bar"),
 					},
 				},
@@ -108,7 +115,7 @@ func TestV3TxnTooManyOps(t *testing.T) {
 			&pb.RequestUnion{
 				Request: &pb.RequestUnion_RequestPut{
 					RequestPut: &pb.PutRequest{
-						Key:   []byte("bar"),
+						Key:   keyf(),
 						Value: []byte("bar"),
 					},
 				},
@@ -132,6 +139,19 @@ func TestV3TxnTooManyOps(t *testing.T) {
 			t.Errorf("#%d: err = %v, want %v", i, err, v3rpc.ErrTooManyOps)
 		}
 	}
+
+	// equal keys fail earlier
+	keyf = func() []byte { return []byte("key") }
+	for i, tt := range tests[1:] {
+		txn := &pb.TxnRequest{}
+		tt(txn)
+		tt(txn)
+
+		_, err := kvc.Txn(context.Background(), txn)
+		if err != v3rpc.ErrDuplicateKey {
+			t.Errorf("#%d: err = %v, want %v", i, err, v3rpc.ErrDuplicateKey)
+		}
+	}
 }
 
 // TestV3PutMissingLease ensures that a Put on a key with a bogus lease fails.