Merge pull request #57 from dismantl/lib-fatal

Remove calls to logrus.Fatal() and panic() in library components

cmd/proxy/main.go

@@ -2,13 +2,14 @@ package main
 
 import (
 	"flag"
+	"os"
+	"strings"
+
 	"github.com/desertbit/grumble"
 	"github.com/hashicorp/yamux"
 	"github.com/nicocha30/ligolo-ng/cmd/proxy/app"
 	"github.com/nicocha30/ligolo-ng/pkg/controller"
 	"github.com/sirupsen/logrus"
-	"os"
-	"strings"
 )
 
 func main() {
@@ -48,7 +49,9 @@ func main() {
 	go proxyController.ListenAndServe()
 
 	// Wait for listener
-	proxyController.WaitForReady()
+	if err := proxyController.WaitForReady(); err != nil {
+		logrus.Fatal(err)
+	}
 
 	// Agent registration goroutine
 	go func() {

pkg/agent/handler.go

@@ -4,16 +4,17 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/nicocha30/ligolo-ng/pkg/agent/neterror"
-	"github.com/nicocha30/ligolo-ng/pkg/agent/smartping"
-	"github.com/nicocha30/ligolo-ng/pkg/protocol"
-	"github.com/nicocha30/ligolo-ng/pkg/relay"
-	"github.com/sirupsen/logrus"
 	"net"
 	"os"
 	"os/user"
 	"syscall"
 	"time"
+
+	"github.com/nicocha30/ligolo-ng/pkg/agent/neterror"
+	"github.com/nicocha30/ligolo-ng/pkg/agent/smartping"
+	"github.com/nicocha30/ligolo-ng/pkg/protocol"
+	"github.com/nicocha30/ligolo-ng/pkg/relay"
+	"github.com/sirupsen/logrus"
 )
 
 var listenerConntrack map[int32]net.Conn
@@ -80,7 +81,8 @@ func NewUDPListener(network string, addr string) (UDPListener, error) {
 func HandleConn(conn net.Conn) {
 	decoder := protocol.NewDecoder(conn)
 	if err := decoder.Decode(); err != nil {
-		panic(err)
+		logrus.Error(err)
+		return
 	}
 
 	e := decoder.Envelope.Payload
@@ -127,7 +129,8 @@ func HandleConn(conn net.Conn) {
 			Type:    protocol.MessageConnectResponse,
 			Payload: connectPacket,
 		}); err != nil {
-			logrus.Fatal(err)
+			logrus.Error(err)
+			return
 		}
 		if connectPacket.Established {
 			relay.StartRelay(targetConn, conn)
@@ -142,7 +145,8 @@ func HandleConn(conn net.Conn) {
 			Type:    protocol.MessageHostPingResponse,
 			Payload: pingResponse,
 		}); err != nil {
-			logrus.Fatal(err)
+			logrus.Error(err)
+			return
 		}
 	case protocol.MessageInfoRequest:
 		var username string
@@ -173,7 +177,8 @@ func HandleConn(conn net.Conn) {
 			Type:    protocol.MessageInfoReply,
 			Payload: infoResponse,
 		}); err != nil {
-			logrus.Fatal(err)
+			logrus.Error(err)
+			return
 		}
 	case protocol.MessageListenerCloseRequest:
 		// Request to close a listener
@@ -325,7 +330,8 @@ func HandleConn(conn net.Conn) {
 			Type:    protocol.MessageListenerSockResponse,
 			Payload: sockResponse,
 		}); err != nil {
-			logrus.Fatal(err)
+			logrus.Error(err)
+			return
 		}
 
 		if sockResponse.Err {

pkg/controller/controller.go

@@ -7,19 +7,21 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
-	"github.com/sirupsen/logrus"
-	"golang.org/x/crypto/acme/autocert"
+	"errors"
 	"math/big"
 	"net"
 	"net/http"
 	"sync"
 	"time"
+
+	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/acme/autocert"
 )
 
 type Controller struct {
 	Network          string
 	Connection       chan net.Conn
-	startchan        chan interface{}
+	startchan        chan error
 	certificateMap   map[string]*tls.Certificate
 	certificateMutex sync.Mutex
 	ControllerConfig
@@ -35,12 +37,11 @@ type ControllerConfig struct {
 }
 
 func New(config ControllerConfig) Controller {
-	return Controller{Network: "tcp", Connection: make(chan net.Conn, 1024), ControllerConfig: config, startchan: make(chan interface{}), certificateMap: make(map[string]*tls.Certificate)}
+	return Controller{Network: "tcp", Connection: make(chan net.Conn, 1024), ControllerConfig: config, startchan: make(chan error), certificateMap: make(map[string]*tls.Certificate)}
 }
 
-func (c *Controller) WaitForReady() {
-	<-c.startchan
-	return
+func (c *Controller) WaitForReady() error {
+	return <-c.startchan
 }
 
 func (c *Controller) ListenAndServe() {
@@ -57,9 +58,18 @@ func (c *Controller) ListenAndServe() {
 			certManager.HostPolicy = autocert.HostWhitelist(c.DomainWhitelist...)
 		}
 		tlsConfig.GetCertificate = certManager.GetCertificate
+
+		// Check if port 80 is available
+		lis, err := net.Listen("tcp", ":http")
+		if err != nil {
+			c.startchan <- errors.New("Port 80 is not available, please make sure it's accessible for Let's Encrypt ACME challenge")
+			return
+		}
+		lis.Close()
+
 		go func() {
 			h := certManager.HTTPHandler(nil)
-			logrus.Fatal(http.ListenAndServe(":http", h))
+			http.ListenAndServe(":http", h)
 		}()
 	} else if c.EnableSelfcert {
 		logrus.Warning("Using automatically generated self-signed certificates (Not recommended)")
@@ -74,7 +84,7 @@ func (c *Controller) ListenAndServe() {
 			c.certificateMutex.Unlock()
 			priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 			if err != nil {
-				logrus.Fatal(err)
+				return nil, err
 			}
 
 			serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
@@ -119,19 +129,23 @@ func (c *Controller) ListenAndServe() {
 	} else if c.Certfile != "" && c.Keyfile != "" {
 		cer, err := tls.LoadX509KeyPair(c.Certfile, c.Keyfile)
 		if err != nil {
-			logrus.WithFields(logrus.Fields{"certfile": c.Certfile, "keyfile": c.Keyfile}).Fatal("Could not load TLS certificate. Please make sure paths are correct or use -autocert or -selfcert options")
+			logrus.WithFields(logrus.Fields{"certfile": c.Certfile, "keyfile": c.Keyfile}).Error("Could not load TLS certificate. Please make sure paths are correct or use -autocert or -selfcert options")
+			c.startchan <- err
+			return
 		}
 		tlsConfig.Certificates = []tls.Certificate{cer}
 	} else {
-		logrus.Fatal("No valid TLS configuration found, please use -certfile/-keyfile, -autocert or -selfcert options")
+		c.startchan <- errors.New("No valid TLS configuration found, please use -certfile/-keyfile, -autocert or -selfcert options")
+		return
 	}
 
 	listener, err := tls.Listen(c.Network, c.Address, &tlsConfig)
 	if err != nil {
-		logrus.Fatal(err)
+		c.startchan <- err
+		return
 	}
 	defer listener.Close()
-	close(c.startchan) // Controller is listening.
+	c.startchan <- nil // Controller is listening.
 	logrus.Infof("Listening on %s", c.Address)
 	for {
 		conn, err := listener.Accept()

pkg/protocol/decoder.go

@@ -42,85 +42,85 @@ func (d *LigoloDecoder) Decode() error {
 	case MessageInfoRequest:
 		p := InfoRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageInfoReply:
 		p := InfoReplyPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageConnectRequest:
 		p := ConnectRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageConnectResponse:
 		p := ConnectResponsePacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageHostPingRequest:
 		p := HostPingRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageHostPingResponse:
 		p := HostPingResponsePacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerRequest:
 		p := ListenerRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerResponse:
 		p := ListenerResponsePacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerBindRequest:
 		p := ListenerBindPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerBindResponse:
 		p := ListenerBindReponse{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerSockRequest:
 		p := ListenerSockRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerSockResponse:
 		p := ListenerSockResponsePacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerCloseRequest:
 		p := ListenerCloseRequestPacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	case MessageListenerCloseResponse:
 		p := ListenerCloseResponsePacket{}
 		if err := gobdecoder.Decode(&p); err != nil {
-			panic(err)
+			return err
 		}
 		d.Envelope.Payload = p
 	default:

pkg/proxy/netstack/icmp.go

@@ -3,6 +3,7 @@ package netstack
 import (
 	"bytes"
 	"errors"
+
 	"github.com/nicocha30/gvisor-ligolo/pkg/buffer"
 	"github.com/nicocha30/gvisor-ligolo/pkg/tcpip"
 	"github.com/nicocha30/gvisor-ligolo/pkg/tcpip/checksum"
@@ -48,7 +49,8 @@ func icmpResponder(s *NetStack) error {
 							continue
 						} else {
 							// This is bad.
-							panic(err)
+							logrus.Error(err)
+							return
 						}
 					}
 
@@ -171,7 +173,7 @@ func ProcessICMP(nstack *stack.Stack, pkt stack.PacketBufferPtr) {
 		replyPkt.TransportProtocolNumber = header.ICMPv4ProtocolNumber
 
 		if err := r.WriteHeaderIncludedPacket(replyPkt); err != nil {
-			panic(err)
+			logrus.Error(err)
 			return
 		}
 	}