diff --git a/daemon/transport/connecter/connecter.go b/daemon/transport/connecter/connecter.go index 3c0f0de..fa772a7 100644 --- a/daemon/transport/connecter/connecter.go +++ b/daemon/transport/connecter/connecter.go @@ -1,12 +1,39 @@ package connecter import ( + "context" "fmt" "github.com/problame/go-streamrpc" "github.com/zrepl/zrepl/config" "github.com/zrepl/zrepl/daemon/streamrpcconfig" + "github.com/zrepl/zrepl/daemon/transport" + "net" + "time" ) + +type HandshakeConnecter struct { + connecter streamrpc.Connecter +} + +func (c HandshakeConnecter) Connect(ctx context.Context) (net.Conn, error) { + conn, err := c.connecter.Connect(ctx) + if err != nil { + return nil, err + } + dl, ok := ctx.Deadline() + if !ok { + dl = time.Now().Add(10 * time.Second) // FIXME constant + } + if err := transport.DoHandshakeCurrentVersion(conn, dl); err != nil { + conn.Close() + return nil, err + } + return conn, nil +} + + + func FromConfig(g *config.Global, in config.ConnectEnum) (*ClientFactory, error) { var ( connecter streamrpc.Connecter @@ -41,6 +68,9 @@ func FromConfig(g *config.Global, in config.ConnectEnum) (*ClientFactory, error) if err := config.Validate(); err != nil { return nil, err } + + connecter = HandshakeConnecter{connecter} + return &ClientFactory{connecter: connecter, config: &config}, nil } diff --git a/daemon/transport/handshake.go b/daemon/transport/handshake.go new file mode 100644 index 0000000..ecfd495 --- /dev/null +++ b/daemon/transport/handshake.go @@ -0,0 +1,136 @@ +package transport + +import ( + "bytes" + "fmt" + "io" + "net" + "strings" + "time" + "unicode/utf8" +) + +type HandshakeMessage struct { + ProtocolVersion int + Extensions []string +} + +func (m *HandshakeMessage) Encode() ([]byte, error) { + if m.ProtocolVersion <= 0 || m.ProtocolVersion > 9999 { + return nil, fmt.Errorf("protocol version must be in [1, 9999]") + } + if len(m.Extensions) >= 9999 { + return nil, fmt.Errorf("protocol only supports [0, 9999] extensions") + } + // EXTENSIONS is a count of subsequent \n separated lines that contain protocol extensions + var extensions strings.Builder + for i, ext := range m.Extensions { + if strings.ContainsAny(ext, "\n") { + return nil, fmt.Errorf("Extension #%d contains forbidden newline character", i) + } + if !utf8.ValidString(ext) { + return nil, fmt.Errorf("Extension #%d is not valid UTF-8", i) + } + extensions.WriteString(ext) + extensions.WriteString("\n") + } + withoutLen := fmt.Sprintf("ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%04d\n%s", + m.ProtocolVersion, len(m.Extensions), extensions.String()) + withLen := fmt.Sprintf("%010d %s", len(withoutLen), withoutLen) + return []byte(withLen), nil +} + +func (m *HandshakeMessage) DecodeReader(r io.Reader, maxLen int) error { + var lenAndSpace [11]byte + if _, err := io.ReadFull(r, lenAndSpace[:]); err != nil { + return err + } + if !utf8.Valid(lenAndSpace[:]) { + return fmt.Errorf("invalid start of handshake message: not valid UTF-8") + } + var followLen int + n, err := fmt.Sscanf(string(lenAndSpace[:]), "%010d ", &followLen) + if n != 1 || err != nil { + return fmt.Errorf("could not parse handshake message length") + } + if followLen > maxLen { + return fmt.Errorf("handshake message length exceeds max length (%d vs %d)", + followLen, maxLen) + } + + var buf bytes.Buffer + _, err = io.Copy(&buf, io.LimitReader(r, int64(followLen))) + if err != nil { + return err + } + + var ( + protoVersion, extensionCount int + ) + n, err = fmt.Fscanf(&buf, "ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%4d\n", + &protoVersion, &extensionCount) + if n != 2 || err != nil { + return fmt.Errorf("could not parse handshake message: %s", err) + } + if protoVersion < 1 { + return fmt.Errorf("invalid protocol version %q", protoVersion) + } + m.ProtocolVersion = protoVersion + + if extensionCount < 0 { + return fmt.Errorf("invalid extension count %q", extensionCount) + } + if extensionCount == 0 { + if buf.Len() != 0 { + return fmt.Errorf("unexpected data trailing after header") + } + m.Extensions = nil + return nil + } + s := buf.String() + if strings.Count(s, "\n") != extensionCount { + return fmt.Errorf("inconsistent extension count: found %d, header says %d", len(m.Extensions), extensionCount) + } + exts := strings.Split(s, "\n") + if exts[len(exts)-1] != "" { + return fmt.Errorf("unexpected data trailing after last extension newline") + } + m.Extensions = exts[0:len(exts)-1] + + return nil +} + +func DoHandshakeCurrentVersion(conn net.Conn, deadline time.Time) error { + // current protocol version is hardcoded here + return DoHandshakeVersion(conn, deadline, 1) +} + +func DoHandshakeVersion(conn net.Conn, deadline time.Time, version int) error { + ours := HandshakeMessage{ + ProtocolVersion: version, + Extensions: nil, + } + hsb, err := ours.Encode() + if err != nil { + return fmt.Errorf("could not encode protocol banner: %s", err) + } + + conn.SetDeadline(deadline) + _, err = io.Copy(conn, bytes.NewBuffer(hsb)) + if err != nil { + return fmt.Errorf("could not send protocol banner: %s", err) + } + + theirs := HandshakeMessage{} + if err := theirs.DecodeReader(conn, 16 * 4096); err != nil { // FIXME constant + return fmt.Errorf("could not decode protocol banner: %s", err) + } + + if theirs.ProtocolVersion != ours.ProtocolVersion { + return fmt.Errorf("protocol versions do not match: ours is %d, theirs is %d", + ours.ProtocolVersion, theirs.ProtocolVersion) + } + // ignore extensions, we don't use them + + return nil +} diff --git a/daemon/transport/handshake_test.go b/daemon/transport/handshake_test.go new file mode 100644 index 0000000..d1c72b4 --- /dev/null +++ b/daemon/transport/handshake_test.go @@ -0,0 +1,119 @@ +package transport + +import ( + "bytes" + "fmt" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/zrepl/zrepl/util/socketpair" + "io" + "strings" + "testing" + "time" +) + +func TestHandshakeMessage_Encode(t *testing.T) { + + msg := HandshakeMessage{ + ProtocolVersion: 2342, + } + + encB, err := msg.Encode() + require.NoError(t, err) + enc := string(encB) + t.Logf("enc: %s", enc) + + + + assert.False(t, strings.ContainsAny(enc[0:10], " ")) + assert.True(t, enc[10] == ' ') + + var ( + headerlen, protoversion, extensionCount int + ) + n, err := fmt.Sscanf(enc, "%010d ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%04d\n", + &headerlen, &protoversion, &extensionCount) + if n != 3 || (err != nil && err != io.EOF) { + t.Fatalf("%v %v", n, err) + } + + assert.Equal(t, 2342, protoversion) + assert.Equal(t, 0, extensionCount) + assert.Equal(t, len(enc)-11, headerlen) + +} + +func TestHandshakeMessage_Encode_InvalidProtocolVersion(t *testing.T) { + + for _, pv := range []int{-1, 0, 10000, 10001} { + t.Logf("testing invalid protocol version = %v", pv) + msg := HandshakeMessage{ + ProtocolVersion: pv, + } + b, err := msg.Encode() + assert.Error(t, err) + assert.Nil(t, b) + } + +} + +func TestHandshakeMessage_DecodeReader(t *testing.T) { + + in := HandshakeMessage{ + 2342, + []string{"foo", "bar 2342"}, + } + + enc, err := in.Encode() + require.NoError(t, err) + + out := HandshakeMessage{} + err = out.DecodeReader(bytes.NewReader([]byte(enc)), 4 * 4096) + assert.NoError(t, err) + assert.Equal(t, 2342, out.ProtocolVersion) + assert.Equal(t, 2, len(out.Extensions)) + assert.Equal(t, "foo", out.Extensions[0]) + assert.Equal(t, "bar 2342", out.Extensions[1]) + +} + +func TestDoHandshakeVersion_ErrorOnDifferentVersions(t *testing.T) { + srv, client, err := socketpair.SocketPair() + if err != nil { + t.Fatal(err) + } + defer srv.Close() + defer client.Close() + + srvErrCh := make(chan error) + go func() { + srvErrCh <- DoHandshakeVersion(srv, time.Now().Add(2*time.Second), 1) + }() + err = DoHandshakeVersion(client, time.Now().Add(2*time.Second), 2) + t.Log(err) + assert.Error(t, err) + assert.True(t, strings.Contains(err.Error(), "version")) + + srvErr := <-srvErrCh + t.Log(srvErr) + assert.Error(t, srvErr) + assert.True(t, strings.Contains(srvErr.Error(), "version")) +} + +func TestDoHandshakeCurrentVersion(t *testing.T) { + srv, client, err := socketpair.SocketPair() + if err != nil { + t.Fatal(err) + } + defer srv.Close() + defer client.Close() + + srvErrCh := make(chan error) + go func() { + srvErrCh <- DoHandshakeVersion(srv, time.Now().Add(2*time.Second), 1) + }() + err = DoHandshakeVersion(client, time.Now().Add(2*time.Second), 1) + assert.Nil(t, err) + assert.Nil(t, <-srvErrCh) + +} diff --git a/daemon/transport/serve/serve.go b/daemon/transport/serve/serve.go index 15e6ba3..c1b3bb1 100644 --- a/daemon/transport/serve/serve.go +++ b/daemon/transport/serve/serve.go @@ -3,12 +3,14 @@ package serve import ( "github.com/pkg/errors" "github.com/zrepl/zrepl/config" + "github.com/zrepl/zrepl/daemon/transport" "net" "github.com/zrepl/zrepl/daemon/streamrpcconfig" "github.com/problame/go-streamrpc" "context" "github.com/zrepl/zrepl/logger" "github.com/zrepl/zrepl/zfs" + "time" ) type contextKey int @@ -71,6 +73,42 @@ type ListenerFactory interface { Listen() (AuthenticatedListener, error) } +type HandshakeListenerFactory struct { + lf ListenerFactory +} + +func (lf HandshakeListenerFactory) Listen() (AuthenticatedListener, error) { + l, err := lf.lf.Listen() + if err != nil { + return nil, err + } + return HandshakeListener{l}, nil +} + +type HandshakeListener struct { + l AuthenticatedListener +} + +func (l HandshakeListener) Addr() (net.Addr) { return l.l.Addr() } + +func (l HandshakeListener) Close() error { return l.l.Close() } + +func (l HandshakeListener) Accept(ctx context.Context) (AuthenticatedConn, error) { + conn, err := l.l.Accept(ctx) + if err != nil { + return nil, err + } + dl, ok := ctx.Deadline() + if !ok { + dl = time.Now().Add(10*time.Second) // FIXME constant + } + if err := transport.DoHandshakeCurrentVersion(conn, dl); err != nil { + conn.Close() + return nil, err + } + return conn, nil +} + func FromConfig(g *config.Global, in config.ServeEnum) (lf ListenerFactory, conf *streamrpc.ConnConfig, _ error) { var ( @@ -100,6 +138,8 @@ func FromConfig(g *config.Global, in config.ServeEnum) (lf ListenerFactory, conf return nil, nil, rpcErr } + lf = HandshakeListenerFactory{lf} + return lf, conf, nil }