From c8bf77c722f40213bd63f71ff8296df2be7c0d2a Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Mon, 25 Jan 2016 21:04:19 -0800
Subject: [PATCH] v3rpc: check max ops in txn

---
 etcdserver/api/v3rpc/error.go |  1 +
 etcdserver/api/v3rpc/key.go   |  8 ++++++
 integration/v3_grpc_test.go   | 53 +++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/etcdserver/api/v3rpc/error.go b/etcdserver/api/v3rpc/error.go
index 35229518e..fb148b83a 100644
--- a/etcdserver/api/v3rpc/error.go
+++ b/etcdserver/api/v3rpc/error.go
@@ -22,6 +22,7 @@ import (
 
 var (
 	ErrEmptyKey      = grpc.Errorf(codes.InvalidArgument, "key is not provided")
+	ErrTooManyOps    = grpc.Errorf(codes.InvalidArgument, "too many operations 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 21daae181..ae1fbcfcb 100644
--- a/etcdserver/api/v3rpc/key.go
+++ b/etcdserver/api/v3rpc/key.go
@@ -27,6 +27,10 @@ import (
 
 var (
 	plog = capnslog.NewPackageLogger("github.com/coreos/etcd/etcdserver/api", "v3rpc")
+
+	// Max operations per txn list. For example, Txn.Success can have at most 128 operations,
+	// and Txn.Failure can have at most 128 operations.
+	MaxOpsPerTxn = 128
 )
 
 type kvServer struct {
@@ -156,6 +160,10 @@ func checkDeleteRequest(r *pb.DeleteRangeRequest) error {
 }
 
 func checkTxnRequest(r *pb.TxnRequest) error {
+	if len(r.Compare) > MaxOpsPerTxn || len(r.Success) > MaxOpsPerTxn || len(r.Failure) > MaxOpsPerTxn {
+		return ErrTooManyOps
+	}
+
 	for _, c := range r.Compare {
 		if len(c.Key) == 0 {
 			return ErrEmptyKey
diff --git a/integration/v3_grpc_test.go b/integration/v3_grpc_test.go
index 0fffa1426..21a64f37b 100644
--- a/integration/v3_grpc_test.go
+++ b/integration/v3_grpc_test.go
@@ -25,6 +25,7 @@ import (
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
 	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
+	"github.com/coreos/etcd/etcdserver/api/v3rpc"
 	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
 	"github.com/coreos/etcd/lease"
 	"github.com/coreos/etcd/storage/storagepb"
@@ -110,6 +111,58 @@ func TestV3PutOverwrite(t *testing.T) {
 	}
 }
 
+func TestV3TxnTooManyOps(t *testing.T) {
+	clus := newClusterGRPC(t, &clusterConfig{size: 3})
+	defer clus.Terminate(t)
+
+	kvc := pb.NewKVClient(clus.RandConn())
+
+	addCompareOps := func(txn *pb.TxnRequest) {
+		txn.Compare = append(txn.Compare,
+			&pb.Compare{
+				Result: pb.Compare_GREATER,
+				Target: pb.Compare_CREATE,
+				Key:    []byte("bar"),
+			})
+	}
+	addSuccessOps := func(txn *pb.TxnRequest) {
+		txn.Success = append(txn.Success,
+			&pb.RequestUnion{
+				RequestPut: &pb.PutRequest{
+					Key:   []byte("bar"),
+					Value: []byte("bar"),
+				},
+			})
+	}
+	addFailureOps := func(txn *pb.TxnRequest) {
+		txn.Failure = append(txn.Failure,
+			&pb.RequestUnion{
+				RequestPut: &pb.PutRequest{
+					Key:   []byte("bar"),
+					Value: []byte("bar"),
+				},
+			})
+	}
+
+	tests := []func(txn *pb.TxnRequest){
+		addCompareOps,
+		addSuccessOps,
+		addFailureOps,
+	}
+
+	for i, tt := range tests {
+		txn := &pb.TxnRequest{}
+		for j := 0; j < v3rpc.MaxOpsPerTxn+1; j++ {
+			tt(txn)
+		}
+
+		_, err := kvc.Txn(context.Background(), txn)
+		if err != v3rpc.ErrTooManyOps {
+			t.Errorf("#%d: err = %v, want %v", i, err, v3rpc.ErrTooManyOps)
+		}
+	}
+}
+
 // TestV3PutMissingLease ensures that a Put on a key with a bogus lease fails.
 func TestV3PutMissingLease(t *testing.T) {
 	clus := newClusterGRPC(t, &clusterConfig{size: 3})