fixing cmd entrypoints

omrikiei · Oct 9, 2020 · 3018cd5 · 3018cd5
1 parent ebd7f6f
commit 3018cd5
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 10 deletions.
diff --git a/cmd/client.go b/cmd/client.go
@@ -45,7 +45,14 @@ ktunnel client --host ktunnel-server.yourcompany.com -s tcp 8000 8001:8432
 			})
 		}()
 
-		err := client.RunClient(ctx, &Host, &Port, Scheme, &Tls, &CaFile, &ServerHostOverride, args)
+		opts := []client.ClientOption{
+			client.WithServer(Host, Port),
+			client.WithTunnels(Scheme, args...),
+			client.WithTLS(CertFile, ServerHostOverride),
+
+		}
+
+		err := client.RunClient(ctx, opts...)
 		if err != nil {
 			log.Fatalf("Failed to run client: %v", err)
 		}

diff --git a/cmd/expose.go b/cmd/expose.go
@@ -82,7 +82,14 @@ ktunnel expose redis 6379
 					log.Fatalf("Failed to run client: %v", err)
 				}
 				prt := int(p)
-				err = client.RunClient(ctx, &Host, &prt, Scheme, &Tls, &CaFile, &ServerHostOverride, args[1:])
+				opts := []client.ClientOption{
+					client.WithServer(Host, prt),
+					client.WithTunnels(Scheme, args[1:]...),
+				}
+				if Tls {
+					opts = append(opts, client.WithTLS(CertFile, ServerHostOverride))
+				}
+				err = client.RunClient(ctx, opts...)
 				if err != nil {
 					log.Fatalf("Failed to run client: %v", err)
 				}

diff --git a/cmd/inject.go b/cmd/inject.go
@@ -89,7 +89,14 @@ ktunnel inject deploymeny mydeployment 3306 6379
 					log.Fatalf("Failed to run client: %v", err)
 				}
 				prt := int(p)
-				err = client.RunClient(ctx, &Host, &prt, Scheme, &Tls, &CaFile, &ServerHostOverride, args[1:])
+				opts := []client.ClientOption{
+					client.WithServer(Host, prt),
+					client.WithTunnels(Scheme, args[1:]...),
+				}
+				if Tls {
+					opts = append(opts, client.WithTLS(CertFile, ServerHostOverride))
+				}
+				err = client.RunClient(ctx, opts...)
 				if err != nil {
 					log.Fatalf("Failed to run client: %v", err)
 				}

diff --git a/cmd/server.go b/cmd/server.go
@@ -40,7 +40,11 @@ ktunnel server -p 8181
 				cancel()
 			})
 		}()
-		err := server.RunServer(ctx, &Port, &Tls, &KeyFile, &CertFile)
+		config := []server.ServerOption{server.WithPort(Port)}
+		if CertFile != "" && KeyFile != "" {
+			config = append(config, server.WithTLS(CertFile, KeyFile))
+		}
+		err := server.RunServer(ctx, config...)
 		if err != nil {
 			log.Fatalf("Error running server: %v", err)
 		}

diff --git a/pkg/client/client.go b/pkg/client/client.go
@@ -3,8 +3,8 @@ package client
 import (
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net"
+	"os"
 	"strings"
 	"time"
 
@@ -150,7 +150,11 @@ func SendData(ctx context.Context, conf *ClientConfig, stream pb.Tunnel_InitTunn
 			session.Lock.Lock()
 			bys := session.Buf.Len()
 			bytes := make([]byte, bys)
-			session.Buf.Read(bytes)
+			_, err := session.Buf.Read(bytes)
+			if err != nil {
+				conf.log.WithError(err).Errorf("failed reading stream from session %v, exiting", err)
+				return
+			}
 
 			conf.log.WithField("session", session.Id).Debugf("read %d from buffer out of %d available", len(bytes), bys)
 
@@ -165,7 +169,7 @@ func SendData(ctx context.Context, conf *ClientConfig, stream pb.Tunnel_InitTunn
 				"session": session.Id,
 				"close":   resp.ShouldClose,
 			}).Debugf("sending %d bytes to server", len(bytes))
-			err := stream.Send(resp)
+			err = stream.Send(resp)
 			if err != nil {
 				conf.log.WithError(err).Errorf("failed sending message to tunnel stream, exiting")
 				return
@@ -250,7 +254,7 @@ func processArgs(opts []ClientOption) (*ClientConfig, error) {
 	// default arguments
 	opt := &ClientConfig{
 		log: &log.Logger{
-			Out: ioutil.Discard,
+			Out: os.Stdout,
 		},
 		scheme: "tcp",
 		TLS:    false,

diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net"
+	"os"
 	"strings"
 
 	"github.com/pkg/errors"
@@ -243,7 +244,7 @@ func RunServer(ctx context.Context, opts ...ServerOption) error {
 	// handle context cancellation, shut down the server
 	go func() {
 		<-ctx.Done()
-		lis.Close()
+		_ = lis.Close()
 	}()
 
 	grpcServer := grpc.NewServer(grpcOpts...)
@@ -257,7 +258,7 @@ func processArgs(opts []ServerOption) (*ServerConfig, error) {
 	opt := &ServerConfig{
 		port: 5000,
 		log: &log.Logger{
-			Out: ioutil.Discard,
+			Out: os.Stdout,
 		},
 		TLS: false,
 	}