Merge pull request #75 from Amnesic-Systems/robustness-improvements

Improve robustness of networking subsystem.

Author: NullHypothesis

cmd/veil-proxy/main.go

@@ -113,6 +113,19 @@ func acceptLoop(ctx context.Context, ln net.Listener, cfg *config.VeilProxy) {
 				log.Printf("Started DNS forwarder at %s.", dns.UDPAddr())
 			}
 
+			// Close vm and tunDev when the context is canceled to
+			// unblock the forwarding goroutines before wg.Wait().
+			stopCh := make(chan struct{})
+			defer close(stopCh)
+			go func() {
+				select {
+				case <-ctx.Done():
+					_ = vm.Close()
+					_ = tunDev.Close()
+				case <-stopCh:
+				}
+			}()
+
 			var wg sync.WaitGroup
 			wg.Add(2)
 			go proxy.VSOCKToTun(vm, tunDev, ch, &wg)

internal/httpx/httpx.go

@@ -90,7 +90,6 @@ func WaitForSvc(
 ) (err error) {
 	defer errs.Wrap(&err, "failed to wait for service")
 
-	start := time.Now()
 	deadline, ok := ctx.Deadline()
 	if !ok {
 		return errors.New("context has no deadline")
@@ -110,7 +109,7 @@ func WaitForSvc(
 			log.Print("Service is ready.")
 			return nil
 		}
-		if time.Since(start) > deadline.Sub(start) {
+		if time.Now().After(deadline) {
 			return errDeadlineExceeded
 		}
 		time.Sleep(10 * time.Millisecond)

internal/net/proxy/proxy.go

@@ -24,18 +24,18 @@ func TunToVSOCK(
 	defer func() { _ = to.Close() }()
 	defer wg.Done()
 	var (
-		err       error
-		pktLenBuf = make([]byte, lenBufSize)
-		pktBuf    = make([]byte, tun.MTU)
+		err     error
+		sendBuf = make([]byte, lenBufSize+tun.MTU)
+		pktBuf  = sendBuf[lenBufSize:]
 	)
 
 	for {
 		// Read a network packet from the tun interface.
 		nr, rerr := from.Read(pktBuf)
 		if nr > 0 {
 			// Forward the network packet to our TCP-over-VSOCK connection.
-			binary.BigEndian.PutUint16(pktLenBuf, uint16(nr))
-			if _, werr := to.Write(append(pktLenBuf, pktBuf[:nr]...)); werr != nil {
+			binary.BigEndian.PutUint16(sendBuf, uint16(nr))
+			if _, werr := to.Write(sendBuf[:lenBufSize+nr]); werr != nil {
 				err = werr
 				break
 			}

internal/net/tun/tun_linux.go

@@ -41,6 +41,7 @@ func setupTun(typ int) (*os.File, error) {
 		return nil, err
 	}
 	if err := configureTun(typ); err != nil {
+		_ = fd.Close()
 		return nil, err
 	}
 

internal/service/service.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"time"
 
 	"github.com/Amnesic-Systems/veil/internal/addr"
 	"github.com/Amnesic-Systems/veil/internal/config"
@@ -140,13 +141,16 @@ func startAllWebSrvs(
 	}(extSrv)
 
 	// Wait until the context is canceled, at which point it's time to stop web
-	// servers.
+	// servers. Use a fresh context for Shutdown so active connections get a
+	// chance to drain rather than being cut off immediately.
 	<-ctx.Done()
 	log.Print("Context cancelled; shutting down veil.")
-	if err := intSrv.Shutdown(ctx); err != nil {
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := intSrv.Shutdown(shutdownCtx); err != nil {
 		log.Printf("Error shutting down internal server: %v", err)
 	}
-	if err := extSrv.Shutdown(ctx); err != nil {
+	if err := extSrv.Shutdown(shutdownCtx); err != nil {
 		log.Printf("Error shutting down external server: %v", err)
 	}
 }