Better classification of ProtocolError

datastax · Jun 18, 2024 · a244c8b · a244c8b
1 parent f60bd28
commit a244c8b
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 2 deletions.
diff --git a/proxy/pkg/zdmproxy/controlconn.go b/proxy/pkg/zdmproxy/controlconn.go
@@ -2,6 +2,7 @@ package zdmproxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"github.com/datastax/go-cassandra-native-protocol/frame"
 	"github.com/datastax/go-cassandra-native-protocol/message"
@@ -378,7 +379,8 @@ func (cc *ControlConn) connAndNegotiateProtoVer(endpoint Endpoint, initialProtoV
 		}
 		newConn := NewCqlConnection(tcpConn, cc.username, cc.password, ccReadTimeout, ccWriteTimeout, cc.conf)
 		err = newConn.InitializeContext(protoVer, ctx)
-		if err != nil && strings.Contains(err.Error(), "Invalid or unsupported protocol version") {
+		var respErr *ResponseError
+		if err != nil && errors.As(err, &respErr) && respErr.IsProtocolError() && strings.Contains(err.Error(), "Invalid or unsupported protocol version") {
 			// unsupported protocol version
 			// protocol renegotiation requires opening a new TCP connection
 			err2 := newConn.Close()

diff --git a/proxy/pkg/zdmproxy/cqlconn.go b/proxy/pkg/zdmproxy/cqlconn.go
@@ -357,6 +357,8 @@ func (c *cqlConn) PerformHandshake(version primitive.ProtocolVersion, ctx contex
 					}
 				}
 			}
+		case *message.ProtocolError:
+			err = &ResponseError{Response: response}
 		default:
 			err = fmt.Errorf("expected AUTHENTICATE or READY, got %v", response.Body.Message)
 		}

diff --git a/proxy/pkg/zdmproxy/response.go b/proxy/pkg/zdmproxy/response.go
@@ -1,6 +1,10 @@
 package zdmproxy
 
-import "github.com/datastax/go-cassandra-native-protocol/frame"
+import (
+	"fmt"
+	"github.com/datastax/go-cassandra-native-protocol/frame"
+	"github.com/datastax/go-cassandra-native-protocol/message"
+)
 
 type Response struct {
 	responseFrame *frame.RawFrame
@@ -37,3 +41,20 @@ func (r *Response) GetStreamId() int16 {
 		return r.requestFrame.Header.StreamId
 	}
 }
+
+type ResponseError struct {
+	Response *frame.Frame
+}
+
+func (pre *ResponseError) Error() string {
+	return fmt.Sprintf("%v", pre.Response.Body.Message)
+}
+
+func (pre *ResponseError) IsProtocolError() bool {
+	switch pre.Response.Body.Message.(type) {
+	case *message.ProtocolError:
+		return true
+	default:
+		return false
+	}
+}