Merge pull request #5 from Amnesic-Systems/bugfixes

Bugfixes and improvements.

nat.go

@@ -12,8 +12,8 @@ const (
 func ToggleNAT(toggle bool) error {
 	var iptablesRules = [][]string{
 		{"nat", "POSTROUTING", "-s", "10.0.0.0/24", "-j", "MASQUERADE"},
-		{"filter", "FORWARD", "-i", TunName, "-s", "10.0.0.0/24", "-j", "ACCEPT"},
-		{"filter", "FORWARD", "-o", TunName, "-d", "10.0.0.0/24", "-j", "ACCEPT"},
+		{"filter", "FORWARD", "-i", tunName, "-s", "10.0.0.0/24", "-j", "ACCEPT"},
+		{"filter", "FORWARD", "-o", tunName, "-d", "10.0.0.0/24", "-j", "ACCEPT"},
 	}
 
 	t, err := iptables.New()

proxy.go

@@ -8,9 +8,9 @@ import (
 )
 
 const (
-	LenBufSize  = 2
-	TunMTU      = 65535 // The maximum-allowed MTU for the tun interface.
-	TunName     = "tun0"
+	lenBufSize  = 2
+	tunMTU      = 65535 // The maximum-allowed MTU for the tun interface.
+	tunName     = "tun0"
 	DefaultPort = 1024
 )
 
@@ -23,8 +23,8 @@ func TunToVsock(from io.Reader, to io.WriteCloser, ch chan error, wg *sync.WaitG
 	defer wg.Done()
 	var (
 		err       error
-		pktLenBuf = make([]byte, LenBufSize)
-		pktBuf    = make([]byte, TunMTU)
+		pktLenBuf = make([]byte, lenBufSize)
+		pktBuf    = make([]byte, tunMTU)
 	)
 
 	for {
@@ -43,7 +43,7 @@ func TunToVsock(from io.Reader, to io.WriteCloser, ch chan error, wg *sync.WaitG
 			break
 		}
 	}
-	ch <- fmt.Errorf("stopped tun-to-vsock forwarding: %v", err)
+	ch <- fmt.Errorf("stopped tun-to-vsock forwarding: %w", err)
 }
 
 // VsockToTun forwards network packets from our TCP-over-VSOCK connection to
@@ -56,8 +56,8 @@ func VsockToTun(from io.Reader, to io.WriteCloser, ch chan error, wg *sync.WaitG
 	var (
 		err       error
 		pktLen    uint16
-		pktLenBuf = make([]byte, LenBufSize)
-		pktBuf    = make([]byte, TunMTU)
+		pktLenBuf = make([]byte, lenBufSize)
+		pktBuf    = make([]byte, tunMTU)
 	)
 
 	for {
@@ -78,5 +78,5 @@ func VsockToTun(from io.Reader, to io.WriteCloser, ch chan error, wg *sync.WaitG
 			break
 		}
 	}
-	ch <- fmt.Errorf("stopped vsock-to-tun forwarding: %v", err)
+	ch <- fmt.Errorf("stopped vsock-to-tun forwarding: %w", err)
 }

proxy_test.go

@@ -2,6 +2,8 @@ package proxy
 
 import (
 	"bytes"
+	"crypto/rand"
+	"errors"
 	"io"
 	"net"
 	"sync"
@@ -15,43 +17,45 @@ func assertEq(t *testing.T, is, should interface{}) {
 	}
 }
 
+// buffer implements io.ReadWriteCloser.
+type buffer struct {
+	*bytes.Buffer
+}
+
+func (b *buffer) Close() error {
+	return nil
+}
+
 func TestAToB(t *testing.T) {
 	var (
-		wg         sync.WaitGroup
-		tun, vsock = net.Pipe()
-		ch         = make(chan error)
-		send       = []byte("hello world")
-		recv       = make([]byte, len(send))
+		err          error
+		wg           sync.WaitGroup
+		ch           = make(chan error)
+		conn1, conn2 = net.Pipe()
+		sendBuf      = make([]byte, tunMTU*2)
+		recvBuf      = &buffer{
+			Buffer: new(bytes.Buffer),
+		}
 	)
 
-	wg.Add(2)
-	go TunToVsock(tun, vsock, ch, &wg)
-	go VsockToTun(vsock, tun, ch, &wg)
-	defer wg.Wait()
-
-	// Read but ignore errors.
-	go func(chan error) {
-		for range ch {
+	// We only expect to see errors containing io.EOF.
+	go func() {
+		for err := range ch {
+			assertEq(t, errors.Is(err, io.EOF), true)
 		}
-	}(ch)
+	}()
 
-	// Echo data back to sender.
-	go func(t *testing.T, expected int) {
-		nw, err := io.Copy(vsock, vsock)
-		assertEq(t, err, nil)
-		assertEq(t, nw, int64(expected))
-	}(t, len(send))
-
-	nw, err := tun.Write(send)
-	assertEq(t, nw, len(send))
-	assertEq(t, err, nil)
-
-	nr, err := tun.Read(recv)
+	// Fill sendBuf with random data.
+	_, err = rand.Read(sendBuf)
 	assertEq(t, err, nil)
-	assertEq(t, nr, len(send))
 
-	err = tun.Close()
-	assertEq(t, err, nil)
-
-	assertEq(t, bytes.Compare(send, recv), 0)
+	wg.Add(2)
+	go TunToVsock(bytes.NewReader(sendBuf), conn1, ch, &wg)
+	go VsockToTun(conn2, recvBuf, ch, &wg)
+	wg.Wait()
+
+	assertEq(t, bytes.Equal(
+		sendBuf,
+		recvBuf.Bytes(),
+	), true)
 }

tun_linux.go

@@ -58,7 +58,7 @@ func createTun() (*os.File, error) {
 	ifr := ifReq{
 		Flags: unix.IFF_TUN | unix.IFF_NO_PI,
 	}
-	copy(ifr.Name[:], TunName)
+	copy(ifr.Name[:], tunName)
 
 	_, _, errno := unix.Syscall(
 		unix.SYS_IOCTL,
@@ -83,7 +83,7 @@ func configureTun(typ int) error {
 		cidrStr = "10.0.0.2/24"
 	}
 
-	link, err := tenus.NewLinkFrom(TunName)
+	link, err := tenus.NewLinkFrom(tunName)
 	if err != nil {
 		return fmt.Errorf("failed to retrieve link: %w", err)
 	}
@@ -94,7 +94,7 @@ func configureTun(typ int) error {
 	if err = link.SetLinkIp(cidr, network); err != nil {
 		return fmt.Errorf("failed to set link address: %w", err)
 	}
-	if err := link.SetLinkMTU(TunMTU); err != nil {
+	if err := link.SetLinkMTU(tunMTU); err != nil {
 		return fmt.Errorf("failed to set link MTU: %w", err)
 	}
 	// Set the enclave's default gateway to the proxy's IP address.