Merge pull request apache#722 from Zariel/fix-head-buf-size

conn: always allocate 9 bytes for the header buf

conn.go

@@ -127,7 +127,7 @@ type Conn struct {
 	timeout time.Duration
 	cfg     *ConnConfig
 
-	headerBuf []byte
+	headerBuf [maxFrameHeaderSize]byte
 
 	streams *streams.IDGenerator
 	mu      sync.RWMutex
@@ -175,11 +175,6 @@ func Connect(host *HostInfo, addr string, cfg *ConnConfig,
 		return nil, err
 	}
 
-	headerSize := 8
-	if cfg.ProtoVersion > protoVersion2 {
-		headerSize = 9
-	}
-
 	c := &Conn{
 		conn:         conn,
 		r:            bufio.NewReader(conn),
@@ -191,7 +186,6 @@ func Connect(host *HostInfo, addr string, cfg *ConnConfig,
 		errorHandler: errorHandler,
 		compressor:   cfg.Compressor,
 		auth:         cfg.Authenticator,
-		headerBuf:    make([]byte, headerSize),
 		quit:         make(chan struct{}),
 		session:      session,
 		streams:      streams.New(cfg.ProtoVersion),
@@ -445,7 +439,7 @@ func (c *Conn) recv() error {
 	}
 
 	// were just reading headers over and over and copy bodies
-	head, err := readHeader(c.r, c.headerBuf)
+	head, err := readHeader(c.r, c.headerBuf[:])
 	if err != nil {
 		return err
 	}

frame.go

@@ -231,6 +231,8 @@ var (
 	ErrFrameTooBig = errors.New("frame length is bigger than the maximum allowed")
 )
 
+const maxFrameHeaderSize = 9
+
 func writeInt(p []byte, n int32) {
 	p[0] = byte(n >> 24)
 	p[1] = byte(n >> 16)
@@ -366,15 +368,15 @@ func readHeader(r io.Reader, p []byte) (head frameHeader, err error) {
 	head.flags = p[1]
 
 	if version > protoVersion2 {
-		if len(p) < 9 {
+		if len(p) != 9 {
 			return frameHeader{}, fmt.Errorf("not enough bytes to read header require 9 got: %d", len(p))
 		}
 
 		head.stream = int(int16(p[2])<<8 | int16(p[3]))
 		head.op = frameOp(p[4])
 		head.length = int(readInt(p[5:]))
 	} else {
-		if len(p) < 8 {
+		if len(p) != 8 {
 			return frameHeader{}, fmt.Errorf("not enough bytes to read header require 8 got: %d", len(p))
 		}