Merge pull request #22 from anfernee/improve

Several improvements: unit test, logging, comments

Author: k8s-ci-robot

cmd/client/main.go

@@ -27,6 +27,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -68,7 +69,6 @@ type GrpcProxyClientOptions struct {
 	proxyHost    string
 	proxyPort    int
 	mode         string
-
 }
 
 func (o *GrpcProxyClientOptions) Flags() *pflag.FlagSet {
@@ -83,6 +83,7 @@ func (o *GrpcProxyClientOptions) Flags() *pflag.FlagSet {
 	flags.StringVar(&o.proxyHost, "proxyHost", o.proxyHost, "The host of the proxy server.")
 	flags.IntVar(&o.proxyPort, "proxyPort", o.proxyPort, "The port the proxy server is listening on.")
 	flags.StringVar(&o.mode, "mode", o.mode, "Mode can be either 'grpc' or 'http-connect'.")
+
 	return flags
 }
 
@@ -130,9 +131,6 @@ func (o *GrpcProxyClientOptions) Validate() error {
 	if o.requestPort > 49151 {
 		return fmt.Errorf("please do not try to use ephemeral port %d for the request server port", o.requestPort)
 	}
-	if o.requestPort < 1024 {
-		return fmt.Errorf("please do not try to use reserved port %d for the request server port", o.requestPort)
-	}
 	if o.proxyPort > 49151 {
 		return fmt.Errorf("please do not try to use ephemeral port %d for the proxy server port", o.proxyPort)
 	}
@@ -182,7 +180,6 @@ func (c *Client) run(o *GrpcProxyClientOptions) error {
 	// Run remote simple http service on server side as
 	// "python -m SimpleHTTPServer"
 
-
 	dialer, err := c.getDialer(o)
 	if err != nil {
 		return fmt.Errorf("failed to get dialer for client, got %v", err)
@@ -202,6 +199,8 @@ func (c *Client) run(o *GrpcProxyClientOptions) error {
 	if err != nil {
 		return fmt.Errorf("failed to send request to client, got %v", err)
 	}
+	defer response.Body.Close() // TODO: proxy server should handle the case where Body isn't closed.
+
 	data, err := ioutil.ReadAll(response.Body)
 	if err != nil {
 		return fmt.Errorf("failed to read response from client, got %v", err)
@@ -233,6 +232,17 @@ func (c *Client) getDialer(o *GrpcProxyClientOptions) (func(ctx context.Context,
 	})
 
 	var proxyConn net.Conn
+
+	// Setup signal handler
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch)
+
+	go func() {
+		<-ch
+		err := proxyConn.Close()
+		klog.Infof("connection closed: %v", err)
+	}()
+
 	switch o.mode {
 	case "grpc":
 		dialOption := grpc.WithTransportCredentials(transportCreds)

pkg/agent/agentclient/client.go

@@ -28,6 +28,8 @@ import (
 	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
 )
 
+// AgentClient runs on the node network side. It connects to proxy server and establishes
+// a stream connection from which it sends and receives network traffic.
 type AgentClient struct {
 	nextConnID  int64
 	connContext map[int64]*connContext
@@ -36,6 +38,17 @@ type AgentClient struct {
 	stream agent.AgentService_ConnectClient
 }
 
+// NewAgentClient creates an AgentClient
+func NewAgentClient(address string) *AgentClient {
+	a := &AgentClient{
+		connContext: make(map[int64]*connContext),
+		address:     address,
+	}
+
+	return a
+}
+
+// connContext tracks a connection from agent to node network.
 type connContext struct {
 	conn      net.Conn
 	cleanFunc func()
@@ -47,15 +60,14 @@ func (c *connContext) cleanup() {
 	c.cleanOnce.Do(c.cleanFunc)
 }
 
-func NewAgentClient(address string) *AgentClient {
-	a := &AgentClient{
-		connContext: make(map[int64]*connContext),
-		address:     address,
-	}
-
-	return a
-}
-
+// Connect connnects to proxy server to establish a gRPC stream,
+// on which the proxied traffic is multiplexed through the stream
+// and piped to the local connection. It register itself as a
+// backend from proxy server, so proxy server will route traffic
+// to this agent.
+//
+// The caller needs to call Serve to start serving proxy requests
+// coming from proxy server.
 func (a *AgentClient) Connect(opts ...grpc.DialOption) error {
 	c, err := grpc.Dial(a.address, opts...)
 	if err != nil {
@@ -72,6 +84,10 @@ func (a *AgentClient) Connect(opts ...grpc.DialOption) error {
 	return nil
 }
 
+// Serve starts to serve proxied requests from proxy server over the
+// gRPC stream. Successful Connect is required before Serve. The
+// The requests include things like opening a connection to a server,
+// streaming data and close the connection.
 func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 	for {
 		select {
@@ -81,8 +97,6 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 		default:
 		}
 
-		klog.Info("waiting packets...")
-
 		pkt, err := a.stream.Recv()
 		if err == io.EOF {
 			klog.Info("received EOF, exit")
@@ -121,6 +135,7 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 				conn:   conn,
 				dataCh: dataCh,
 				cleanFunc: func() {
+					klog.Infof("close connection(id=%d)", connID)
 					resp := &agent.Packet{
 						Type:    agent.PacketType_CLOSE_RSP,
 						Payload: &agent.Packet_CloseResponse{CloseResponse: &agent.CloseResponse{}},
@@ -201,11 +216,13 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
 
 	for {
 		n, err := conn.Read(buf[:])
+		klog.Infof("received %d bytes from proxy server", n)
 
 		if err == io.EOF {
 			klog.Info("connection EOF")
 			return
 		} else if err != nil {
+			// Normal when receive a CLOSE_REQ
 			klog.Warningf("connection read error: %v", err)
 			return
 		} else {

pkg/agent/agentserver/server.go

@@ -39,14 +39,16 @@ func (c *ProxyClientConnection) send(pkt *agent.Packet) error {
 		stream := c.Grpc
 		return stream.Send(pkt)
 	} else if c.Mode == "http-connect" {
-		if pkt.Type != agent.PacketType_DATA {
-			return nil
+		if pkt.Type == agent.PacketType_CLOSE_RSP {
+			return c.Http.Close()
+		} else if pkt.Type == agent.PacketType_DATA {
+			_, err := c.Http.Write(pkt.GetData().Data)
+			return err
+		} else {
+			return fmt.Errorf("attempt to send via unrecognized connection type %v", pkt.Type)
 		}
-		writer := c.Http
-		_, err := writer.Write(pkt.GetData().Data)
-		return err
 	} else {
-		return fmt.Errorf("attempt to send via unrecognized connection type %q", c.Mode)
+		return fmt.Errorf("attempt to send via unrecognized connection mode %q", c.Mode)
 	}
 }
 
@@ -84,6 +86,7 @@ func (s *ProxyServer) Proxy(stream agent.ProxyService_ProxyServer) error {
 		close(recvCh)
 	}()
 
+	// Start goroutine to receive packets from frontend and push to recvCh
 	go func() {
 		for {
 			in, err := stream.Recv()
@@ -92,7 +95,8 @@ func (s *ProxyServer) Proxy(stream agent.ProxyService_ProxyServer) error {
 				return
 			}
 			if err != nil {
-				klog.Warningf("stream read error: %v", err)
+				klog.Warningf(">>> Stream read from frontend error: %v", err)
+				close(stopCh)
 				return
 			}
 
@@ -104,57 +108,82 @@ func (s *ProxyServer) Proxy(stream agent.ProxyService_ProxyServer) error {
 }
 
 func (s *ProxyServer) serveRecvFrontend(stream agent.ProxyService_ProxyServer, recvCh <-chan *agent.Packet) {
-	klog.Info("start serve recv ...")
+	klog.Info("start serving frontend stream")
+
+	var firstConnID int64
+
 	for pkt := range recvCh {
 		switch pkt.Type {
 		case agent.PacketType_DIAL_REQ:
-			klog.Info("received DIAL_REQ")
+			klog.Info(">>> Received DIAL_REQ")
 			if s.Backend == nil {
-				klog.Info("no backend found; drop")
+				klog.Info(">>> No backend found; drop")
 				continue
 			}
 
 			if err := s.Backend.Send(pkt); err != nil {
-				klog.Warningf("send packet to Backend failed: %v", err)
+				klog.Warningf(">>> DIAL_REQ to Backend failed: %v", err)
 			}
 			s.PendingDial[pkt.GetDialRequest().Random] = &ProxyClientConnection{
 				Mode:      "grpc",
 				Grpc:      stream,
 				connected: make(chan struct{}),
 			}
-			klog.Info("DIAL_REQ sent to backend") // got this. but backend didn't receive anything.
+			klog.Info(">>> DIAL_REQ sent to backend") // got this. but backend didn't receive anything.
 
 		case agent.PacketType_CLOSE_REQ:
-			klog.Infof("received CLOSE_REQ(id=%d)", pkt.GetCloseRequest().ConnectID)
+			klog.Infof(">>> Received CLOSE_REQ(id=%d)", pkt.GetCloseRequest().ConnectID)
 			if s.Backend == nil {
-				klog.Info("no backend found; drop")
+				klog.Info(">>> No backend found; drop")
 				continue
 			}
 
 			if err := s.Backend.Send(pkt); err != nil {
-				klog.Warningf("send packet to Backend failed: %v", err)
+				klog.Warningf(">>> CLOSE_REQ to Backend failed: %v", err)
 			}
 			klog.Info("CLOSE_REQ sent to backend")
 
 		case agent.PacketType_DATA:
-			klog.Infof("received DATA(id=%d)", pkt.GetData().ConnectID)
+			connID := pkt.GetData().ConnectID
+			klog.Infof(">>> Received DATA(id=%d)", connID)
+			if firstConnID == 0 {
+				firstConnID = connID
+			} else if firstConnID != connID {
+				klog.Warningf(">>> Data(id=%d) doesn't match first connection id %d", firstConnID, connID)
+			}
+
 			if s.Backend == nil {
-				klog.Info("no backend found; drop")
+				klog.Info(">>> No backend found; drop")
 				continue
 			}
 
 			if err := s.Backend.Send(pkt); err != nil {
-				klog.Warningf("send packet to Backend failed: %v", err)
+				klog.Warningf(">>> DATA to Backend failed: %v", err)
 			}
-			klog.Info("DATA sent to backend")
+			klog.Info(">>> DATA sent to backend")
 
 		default:
-			klog.Infof("Ignore %v packet coming from frontend", pkt.Type)
+			klog.Infof(">>> Ignore %v packet coming from frontend", pkt.Type)
+		}
+	}
+
+	klog.Infof(">>> Close streaming (id=%d)", firstConnID)
+
+	pkt := &agent.Packet{
+		Type: agent.PacketType_CLOSE_REQ,
+		Payload: &agent.Packet_CloseRequest{
+			CloseRequest: &agent.CloseRequest{
+				ConnectID: firstConnID,
+			},
+		},
+	}
+	if s.Backend != nil {
+		if err := s.Backend.Send(pkt); err != nil {
+			klog.Warningf(">>> CLOSE_REQ to Backend failed: %v", err)
 		}
 	}
 }
 
-// Ignored now
 func (s *ProxyServer) serveSend(stream agent.ProxyService_ProxyServer, sendCh <-chan *agent.Packet) {
 	klog.Info("start serve send ...")
 	for pkt := range sendCh {
@@ -206,40 +235,56 @@ func (s *ProxyServer) Connect(stream agent.AgentService_ConnectServer) error {
 
 // route the packet back to the correct client
 func (s *ProxyServer) serveRecvBackend(stream agent.AgentService_ConnectServer, recvCh <-chan *agent.Packet) {
+	var firstConnID int64
+
 	for pkt := range recvCh {
 		switch pkt.Type {
 		case agent.PacketType_DIAL_RSP:
 			resp := pkt.GetDialResponse()
-			klog.Warningf("Received dial response for %d, connectID is %d", resp.Random, resp.ConnectID)
+			firstConnID = resp.ConnectID
+			klog.Infof("<<< Received DIAL_RSP(rand=%d, id=%d)", resp.Random, resp.ConnectID)
+
 			if client, ok := s.PendingDial[resp.Random]; !ok {
-				klog.Warning("DialResp not recognized; dropped")
+				klog.Warning("<<< DialResp not recognized; dropped")
 			} else {
 				err := client.send(pkt)
 				delete(s.PendingDial, resp.Random)
 				if err != nil {
-					klog.Warningf("dial response send to client stream error: %v", err)
+					klog.Warningf("<<< DIAL_RSP send to client stream error: %v", err)
 				} else {
 					client.connectID = resp.ConnectID
 					s.Frontends[resp.ConnectID] = client
 					close(client.connected)
 				}
 			}
+
 		case agent.PacketType_DATA:
 			resp := pkt.GetData()
+			klog.Infof("<<< Received DATA(id=%d)", resp.ConnectID)
 			if client, ok := s.Frontends[resp.ConnectID]; ok {
 				if err := client.send(pkt); err != nil {
-					klog.Warningf("data send to client stream error: %v", err)
+					klog.Warningf("<<< DATA send to client stream error: %v", err)
+				} else {
+					klog.Infof("<<< DATA sent to frontend")
 				}
 			}
+
 		case agent.PacketType_CLOSE_RSP:
 			resp := pkt.GetCloseResponse()
+			klog.Infof("<<< Received CLOSE_RSP(id=%d)", resp.ConnectID)
 			if client, ok := s.Frontends[resp.ConnectID]; ok {
 				if err := client.send(pkt); err != nil {
-					klog.Warningf("close response send to client stream error: %v", err)
+					// Normal when frontend closes it.
+					klog.Warningf("<<< CLOSE_RSP send to client stream error: %v", err)
+				} else {
+					klog.Infof("<<< CLOSE_RSP sent to frontend")
 				}
 			}
+
 		default:
-			klog.Warningf("unrecognized packet %+v", pkt)
+			klog.Warningf("<<< Unrecognized packet %+v", pkt)
 		}
 	}
-}
+
+	klog.Infof("<<< Close streaming (id=%d)", firstConnID)
+}

pkg/agent/agentserver/tunnel.go

@@ -99,16 +99,18 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	klog.Infof("Starting proxy to %q", r.Host)
 	pkt := make([]byte, 1<<12)
+
 	for {
 		n, err := conn.Read(pkt[:])
 		if err == io.EOF {
-			// TODO: Close remote..
+			klog.Warningf("EOF from %v", r.Host)
 			break
 		}
 		if err != nil {
 			klog.Errorf("Received error on connection %v", err)
 			break
 		}
+
 		packet := &agent.Packet{
 			Type: agent.PacketType_DATA,
 			Payload: &agent.Packet_Data{