Handshaker component, part 2: Handshaker implementation

security/s2a/internal/handshaker/handshaker.go

@@ -21,12 +21,21 @@ package handshaker
 import (
 	"context"
 	"errors"
+	"fmt"
+	"io"
 	"net"
 
-	"google.golang.org/grpc"
+	grpc "google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/security/s2a/internal/authinfo"
 	s2apb "google.golang.org/grpc/security/s2a/internal/proto"
 )
 
+var (
+	appProtocols = []string{"grpc"}
+	frameLimit   = 1024 * 128
+)
+
 // ClientHandshakerOptions contains the options needed to configure the S2A
 // handshaker service on the client-side.
 type ClientHandshakerOptions struct {
@@ -112,30 +121,164 @@ func newServerHandshaker(stream s2apb.S2AService_SetUpSessionClient, c net.Conn,
 	}
 }
 
-// ClientHandshake performs a client-side TLS handshake using S2A. When complete,
-// it returns a TLS connection.
-func (h *s2aHandshaker) ClientHandshake(ctx context.Context) (net.Conn, error) {
-	return nil, errors.New("Method unimplemented")
+// ClientHandshake performs a client-side TLS handshake using the S2A handshaker
+// service. When complete, returns a secure TLS connection.
+func (h *s2aHandshaker) ClientHandshake(ctx context.Context) (net.Conn, *authinfo.S2AAuthInfo, error) {
+	if h.clientOpts == nil {
+		return nil, nil, errors.New("only handshakers created using NewClientHandshaker can perform a client handshaker")
+	}
+
+	// Create target identities from service account list.
+	req := &s2apb.SessionReq{
+		ReqOneof: &s2apb.SessionReq_ClientStart{
+			ClientStart: &s2apb.ClientSessionStartReq{
+				ApplicationProtocols: appProtocols,
+				MinTlsVersion:        h.clientOpts.MinTLSVersion,
+				MaxTlsVersion:        h.clientOpts.MaxTLSVersion,
+				TlsCiphersuites:      h.clientOpts.TLSCiphersuites,
+				TargetIdentities:     h.clientOpts.TargetIdentities,
+				LocalIdentity:        h.clientOpts.LocalIdentity,
+				TargetName:           h.clientOpts.TargetName,
+			},
+		},
+	}
+
+	conn, result, err := h.setUpSession(req)
+	if err != nil {
+		return nil, nil, err
+	}
+	authInfo, err := authinfo.NewS2AAuthInfo(result)
+	if err != nil {
+		return nil, nil, err
+	}
+	return conn, authInfo, nil
 }
 
 // ServerHandshake performs a server-side TLS handshake using the S2A handshaker
 // service. When complete, returns a secure TLS connection.
-func (h *s2aHandshaker) ServerHandshake(ctx context.Context) (net.Conn, error) {
-	return nil, errors.New("Method unimplemented")
+func (h *s2aHandshaker) ServerHandshake(ctx context.Context) (net.Conn, *authinfo.S2AAuthInfo, error) {
+
+	if h.serverOpts == nil {
+		return nil, nil, errors.New("only handshakers created using NewServerHandshaker can perform a server handshaker")
+	}
+
+	p := make([]byte, 64*1024) // temp length?
+	n, err := h.conn.Read(p)
+	if err != nil {
+		return nil, nil, err
+	}
+	req := &s2apb.SessionReq{
+		ReqOneof: &s2apb.SessionReq_ServerStart{
+			ServerStart: &s2apb.ServerSessionStartReq{
+				ApplicationProtocols: appProtocols,
+				MinTlsVersion:        h.serverOpts.MinTLSVersion,
+				MaxTlsVersion:        h.serverOpts.MaxTLSVersion,
+				TlsCiphersuites:      h.serverOpts.TLSCiphersuites,
+				LocalIdentities:      h.serverOpts.LocalIdentities,
+				InBytes:              p[:n],
+			},
+		},
+	}
+
+	conn, result, err := h.setUpSession(req)
+	if err != nil {
+		return nil, nil, err
+	}
+	authInfo, err := authinfo.NewS2AAuthInfo(result)
+	if err != nil {
+		return nil, nil, err
+	}
+	return conn, authInfo, nil
 }
 
 func (h *s2aHandshaker) setUpSession(req *s2apb.SessionReq) (net.Conn, *s2apb.SessionResult, error) {
-	return nil, nil, errors.New("Method unimplemented")
+	resp, err := h.accessHandshakerService(req)
+	if err != nil {
+		return nil, nil, err
+	}
+	// Check if the returned status is an error.
+	if resp.GetStatus() != nil {
+		if got, want := resp.GetStatus().Code, uint32(codes.OK); got != want {
+			return nil, nil, fmt.Errorf("%v", resp.GetStatus().Details)
+		}
+	}
+
+	var extra []byte
+	if req.GetServerStart() != nil {
+		if resp.GetBytesConsumed() > uint32(len(req.GetServerStart().GetInBytes())) {
+			return nil, nil, errors.New("handshaker service consumed bytes value is out-of-bound")
+		}
+		extra = req.GetServerStart().GetInBytes()[resp.GetBytesConsumed():]
+	}
+	result, extra, err := h.processUntilDone(resp, extra)
+	if err != nil {
+		return nil, nil, err
+	}
+	// TODO: implemented record protocol & new Conn
+	return h.conn, result, nil
 }
 
 func (h *s2aHandshaker) accessHandshakerService(req *s2apb.SessionReq) (*s2apb.SessionResp, error) {
-	return nil, errors.New("Method unimplemented")
+	if err := h.stream.Send(req); err != nil {
+		return nil, err
+	}
+	resp, err := h.stream.Recv()
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
 }
 
+// processUntilDone processes the handshake until the handshaker service returns
+// the results.
 func (h *s2aHandshaker) processUntilDone(resp *s2apb.SessionResp, extra []byte) (*s2apb.SessionResult, []byte, error) {
-	return nil, nil, errors.New("Method unimplemented")
+	for {
+		if len(resp.OutFrames) > 0 {
+			if _, err := h.conn.Write(resp.OutFrames); err != nil {
+				return nil, nil, err
+			}
+		}
+		if resp.Result != nil {
+			return resp.Result, extra, nil
+		}
+		buf := make([]byte, frameLimit)
+		n, err := h.conn.Read(buf)
+		if err != nil && err != io.EOF {
+			return nil, nil, err
+		}
+		// If there is nothing to send to the handshaker service, and
+		// nothing is received from the peer, then we are stuck.
+		// This covers the case when the peer is not responding. Note
+		// that handshaker service connection issues are caught in
+		// accessHandshakerService before we even get here.
+		if len(resp.OutFrames) == 0 && n == 0 {
+			return nil, nil, errors.New("peer server is not responding and re-connection should be attempted")
+		}
+		// Append extra bytes from the previous interaction with the
+		// handshaker service with the current buffer read from conn.
+		p := append(extra, buf[:n]...)
+		// From here on, p and extra point to the same slice.
+		resp, err = h.accessHandshakerService(&s2apb.SessionReq{
+			ReqOneof: &s2apb.SessionReq_Next{
+				Next: &s2apb.SessionNextReq{
+					InBytes: p,
+				},
+			},
+		})
+		if err != nil {
+			return nil, nil, err
+		}
+		// Set extra based on handshaker service response.
+		if resp.GetBytesConsumed() > uint32(len(p)) {
+			return nil, nil, errors.New("handshaker service consumed bytes value is out-of-bound")
+		}
+		extra = p[resp.GetBytesConsumed():]
+	}
 }
 
+// Close shuts down the handshaker and the stream to the S2A handshaker service
+// when the handshake is complete. It should be called when the caller obtains
+// the secure connection at the end of the handshake; otherwise it is a no-op.
 func (h *s2aHandshaker) Close() {
-	// Method is unimplemented.
+	h.stream.CloseSend()
 }