make v5 works as well

Author: worryg0d

conn.go

@@ -410,7 +410,7 @@ func (s *startupCoordinator) options(ctx context.Context, startupCompleted *atom
 
 	supported, ok := frame.(*supportedFrame)
 	if !ok {
-		return NewErrProtocol("Unknown type of response to startup frame: %T", frame)
+		return NewErrProtocol("Unknown type of response to startup frame: %T (frame=%s)", frame, frame.String())
 	}
 
 	return s.startup(ctx, supported.supported, startupCompleted)

conn_test.go

@@ -1172,9 +1172,26 @@ func (srv *TestServer) serve() {
 		}
 
 		go func(conn net.Conn) {
+			var startupCompleted bool
+			var useProtoV5 bool
+
 			defer conn.Close()
 			for !srv.isClosed() {
-				framer, err := srv.readFrame(conn)
+				var reader io.Reader = conn
+
+				if useProtoV5 && startupCompleted {
+					frame, _, err := readUncompressedSegment(conn)
+					if err != nil {
+						if errors.Is(err, io.EOF) {
+							return
+						}
+						srv.errorLocked(err)
+						return
+					}
+					reader = bytes.NewReader(frame)
+				}
+
+				framer, err := srv.readFrame(reader)
 				if err != nil {
 					if err == io.EOF {
 						return
@@ -1187,7 +1204,7 @@ func (srv *TestServer) serve() {
 					srv.onRecv(framer)
 				}
 
-				go srv.process(conn, framer)
+				srv.process(conn, framer, &useProtoV5, &startupCompleted)
 			}
 		}(conn)
 	}
@@ -1225,7 +1242,7 @@ func (srv *TestServer) errorLocked(err interface{}) {
 	srv.t.Error(err)
 }
 
-func (srv *TestServer) process(conn net.Conn, reqFrame *framer) {
+func (srv *TestServer) process(conn net.Conn, reqFrame *framer, useProtoV5, startupCompleted *bool) {
 	head := reqFrame.header
 	if head == nil {
 		srv.errorLocked("process frame with a nil header")
@@ -1238,14 +1255,8 @@ func (srv *TestServer) process(conn net.Conn, reqFrame *framer) {
 			srv.errorLocked(err)
 			return
 		}
-		respFrame.buf[0] |= 0x80
-		if err := respFrame.finish(); err != nil {
-			srv.errorLocked(err)
-		}
-		if err := respFrame.writeTo(conn); err != nil {
-			srv.errorLocked(err)
-		}
-		return
+		// Dont like this but...
+		goto finish
 	}
 
 	switch head.op {
@@ -1437,34 +1448,54 @@ func (srv *TestServer) process(conn net.Conn, reqFrame *framer) {
 		respFrame.writeString("not supported")
 	}
 
-	respFrame.buf[0] = srv.protocol | 0x80
+finish:
+
+	respFrame.buf[0] |= 0x80
 
 	if err := respFrame.finish(); err != nil {
 		srv.errorLocked(err)
 	}
 
-	if err := respFrame.writeTo(conn); err != nil {
-		srv.errorLocked(err)
+	if *useProtoV5 && *startupCompleted {
+		segment, err := newUncompressedSegment(respFrame.buf, true)
+		if err == nil {
+			_, err = conn.Write(segment)
+		}
+		if err != nil {
+			srv.errorLocked(err)
+			return
+		}
+	} else {
+		if err := respFrame.writeTo(conn); err != nil {
+			srv.errorLocked(err)
+		}
+
+		if reqFrame.header.op == opStartup {
+			*startupCompleted = true
+			if head.version == protoVersion5 {
+				*useProtoV5 = true
+			}
+		}
 	}
 }
 
-func (srv *TestServer) readFrame(conn net.Conn) (*framer, error) {
+func (srv *TestServer) readFrame(reader io.Reader) (*framer, error) {
 	buf := make([]byte, srv.headerSize)
-	head, err := readHeader(conn, buf)
+	head, err := readHeader(reader, buf)
 	if err != nil {
 		return nil, err
 	}
 	framer := newFramer(nil, srv.protocol, GlobalTypes)
 
-	err = framer.readFrame(conn, &head)
+	err = framer.readFrame(reader, &head)
 	if err != nil {
 		return nil, err
 	}
 
 	// should be a request frame
 	if head.version.response() {
 		return nil, fmt.Errorf("expected to read a request frame got version: %v", head.version)
-	} else if head.version.version() != srv.protocol {
+	} else if !srv.dontFailOnProtocolMismatch && head.version.version() != srv.protocol {
 		return nil, fmt.Errorf("expected to read protocol version 0x%x got 0x%x", srv.protocol, head.version.version())
 	}
 

control.go

@@ -201,8 +201,8 @@ func shuffleHosts(hosts []*HostInfo) []*HostInfo {
 	return shuffled
 }
 
-const highestProtocolVersionSupported = 5
-const lowestProtocolVersionSupported = 3
+const highestProtocolVersionSupported = protoVersion5
+const lowestProtocolVersionSupported = protoVersion3
 
 func (c *controlConn) discoverProtocol(hosts []*HostInfo) (int, error) {
 	hosts = shuffleHosts(hosts)