Fix socks server and TCP service mapping

Signed-off-by: Vasyl Gello <vasek.gello@gmail.com>

cmd/yggstack/main.go

@@ -1,12 +1,10 @@
 package main
 
 import (
-	"crypto/ed25519"
-	"encoding/hex"
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
-	"net"
 	"os"
 	"os/signal"
 	"syscall"
@@ -14,13 +12,10 @@ import (
 	"github.com/gologme/log"
 	gsyslog "github.com/hashicorp/go-syslog"
 	"github.com/hjson/hjson-go"
-	"github.com/things-go/go-socks5"
 
-	"github.com/yggdrasil-network/yggdrasil-go/cmd/yggstack/types"
-	"github.com/yggdrasil-network/yggdrasil-go/contrib/netstack"
-	"github.com/yggdrasil-network/yggdrasil-go/src/address"
 	"github.com/yggdrasil-network/yggdrasil-go/src/config"
 	"github.com/yggdrasil-network/yggdrasil-go/src/setup"
+	"github.com/yggdrasil-network/yggdrasil-go/src/types"
 
 	"github.com/yggdrasil-network/yggdrasil-go/src/version"
 
@@ -33,8 +28,12 @@ func main() {
 	socks := flag.String("socks", "", "address to listen on for SOCKS, i.e. :1080")
 	nameserver := flag.String("nameserver", "", "the Yggdrasil IPv6 address to use as a DNS server for SOCKS")
 	flag.Var(&expose, "exposetcp", "TCP ports to expose to the network, e.g. 22, 2022:22, 22:192.168.1.1:2022")
+
 	args := setup.ParseArguments()
 
+	// Catch interrupts from the operating system to exit gracefully.
+	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+
 	// Create a new logger that logs output to stdout.
 	var logger *log.Logger
 	switch args.LogTo {
@@ -104,95 +103,21 @@ func main() {
 	n := setup.NewNode(cfg, logger)
 	n.SetLogLevel(args.LogLevel)
 
-	// Have we been asked for the node address yet? If so, print it and then stop.
-	getNodeKey := func() ed25519.PublicKey {
-		if pubkey, err := hex.DecodeString(cfg.PrivateKey); err == nil {
-			return ed25519.PrivateKey(pubkey).Public().(ed25519.PublicKey)
-		}
-		return nil
-	}
-	switch {
-	case args.GetAddr:
-		if key := getNodeKey(); key != nil {
-			addr := address.AddrForKey(key)
-			ip := net.IP(addr[:])
-			fmt.Println(ip.String())
-		}
-		return
-	case args.GetSubnet:
-		if key := getNodeKey(); key != nil {
-			snet := address.SubnetForKey(key)
-			ipnet := net.IPNet{
-				IP:   append(snet[:], 0, 0, 0, 0, 0, 0, 0, 0),
-				Mask: net.CIDRMask(len(snet)*8, 128),
-			}
-			fmt.Println(ipnet.String())
-		}
-		return
-	default:
-	}
-
-	// Now start Yggdrasil - this starts the DHT, router, switch and other core
+	// Now start Yggdrasil - this starts the router, switch and other core
 	// components needed for Yggdrasil to operate
 	if err = n.Run(args); err != nil {
 		logger.Fatalln(err)
 	}
 
-	// Make some nice output that tells us what our IPv6 address and subnet are.
-	// This is just logged to stdout for the user.
-	address := n.Address()
-	subnet := n.Subnet()
-	public := n.GetSelf().Key
-	publicstr := hex.EncodeToString(public[:])
-	logger.Infof("Your public key is %s", publicstr)
-	logger.Infof("Your IPv6 address is %s", address.String())
-	logger.Infof("Your IPv6 subnet is %s", subnet.String())
-	logger.Infof("Your Yggstack resolver name is %s%s", publicstr, types.NameMappingSuffix)
-
-	s, err := netstack.CreateYggdrasilNetstack(&n.Core)
+	// Create Yggdrasil netstack
+	err = n.SetupNetstack(socks, nameserver, &expose)
 	if err != nil {
 		logger.Fatalln(err)
 	}
 
-	if *socks != "" {
-		resolver := types.NewNameResolver(s, *nameserver)
-		server := socks5.NewServer(
-			socks5.WithDial(s.DialContext),
-			socks5.WithResolver(resolver),
-		)
-		go server.ListenAndServe("tcp", *socks) // nolint:errcheck
-	}
-
-	for _, mapping := range expose {
-		go func(mapping types.TCPMapping) {
-			listener, err := s.ListenTCP(mapping.Listen)
-			if err != nil {
-				panic(err)
-			}
-			logger.Infof("Mapping Yggdrasil port %d to %s", mapping.Listen.Port, mapping.Mapped)
-			for {
-				c, err := listener.Accept()
-				if err != nil {
-					panic(err)
-				}
-				r, err := net.DialTCP("tcp", nil, mapping.Mapped)
-				if err != nil {
-					logger.Errorf("Failed to connect to %s: %s", mapping.Mapped, err)
-					_ = c.Close()
-					continue
-				}
-				types.ProxyTCP(n.MTU(), c, r)
-			}
-		}(mapping)
-	}
-
-	term := make(chan os.Signal, 1)
-	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
-
-	select {
-	case <-n.Done():
-	case <-term:
-	}
+	// Block until we are told to shut down.
+	<-ctx.Done()
 
+	// Shut down the node.
 	n.Close()
 }

contrib/netstack/netstack.go

@@ -41,9 +41,14 @@ func CreateYggdrasilNetstack(ygg *core.Core) (*YggdrasilNetstack, error) {
 }
 
 func convertToFullAddr(ip net.IP, port int) (tcpip.FullAddress, tcpip.NetworkProtocolNumber, error) {
+	addr := tcpip.Address{}
+	ip16 := ip.To16()
+	if ip16 != nil {
+		addr = tcpip.AddrFromSlice(ip16)
+	}
 	return tcpip.FullAddress{
 		NIC:  1,
-		Addr: tcpip.Address(ip),
+		Addr: addr,
 		Port: uint16(port),
 	}, ipv6.ProtocolNumber, nil
 }

contrib/netstack/yggdrasil.go

@@ -43,11 +43,9 @@ func (s *YggdrasilNetstack) NewYggdrasilNIC(ygg *core.Core) tcpip.Error {
 				break
 			}
 			pkb := stack.NewPacketBuffer(stack.PacketBufferOptions{
-				Data: buffer.NewVectorisedView(rx, []buffer.View{
-					buffer.NewViewFromBytes(nic.readBuf[:rx]),
-				}),
+				Payload: buffer.MakeWithData(nic.readBuf[:rx]),
 			})
-			nic.dispatcher.DeliverNetworkPacket("", "", ipv6.ProtocolNumber, pkb)
+			nic.dispatcher.DeliverNetworkPacket(ipv6.ProtocolNumber, pkb)
 		}
 	}()
 	_, snet, err := net.ParseCIDR("0200::/7")
@@ -55,8 +53,8 @@ func (s *YggdrasilNetstack) NewYggdrasilNIC(ygg *core.Core) tcpip.Error {
 		return &tcpip.ErrBadAddress{}
 	}
 	subnet, err := tcpip.NewSubnet(
-		tcpip.Address(string(snet.IP)),
-		tcpip.AddressMask(string(snet.Mask)),
+		tcpip.AddrFromSlice(snet.IP.To16()),
+		tcpip.MaskFrom(string(snet.Mask)),
 	)
 	if err != nil {
 		return &tcpip.ErrBadAddress{}
@@ -71,7 +69,7 @@ func (s *YggdrasilNetstack) NewYggdrasilNIC(ygg *core.Core) tcpip.Error {
 			1,
 			tcpip.ProtocolAddress{
 				Protocol:          ipv6.ProtocolNumber,
-				AddressWithPrefix: tcpip.Address(ip).WithPrefix(),
+				AddressWithPrefix: tcpip.AddrFromSlice(ip.To16()).WithPrefix(),
 			},
 			stack.AddressProperties{},
 		); err != nil {
@@ -95,31 +93,25 @@ func (*YggdrasilNIC) LinkAddress() tcpip.LinkAddress { return "" }
 
 func (*YggdrasilNIC) Wait() {}
 
-func (e *YggdrasilNIC) WritePacket(
-	_ stack.RouteInfo,
-	_ tcpip.NetworkProtocolNumber,
-	pkt *stack.PacketBuffer,
-) tcpip.Error {
-	vv := buffer.NewVectorisedView(pkt.Size(), pkt.Views())
-	n, err := vv.Read(e.writeBuf)
-	if err != nil {
-		log.Println(err)
-		return &tcpip.ErrAborted{}
-	}
-	_, err = e.ipv6rwc.Write(e.writeBuf[:n])
-	if err != nil {
-		log.Println(err)
-		return &tcpip.ErrAborted{}
-	}
-	return nil
-}
-
 func (e *YggdrasilNIC) WritePackets(
-	stack.RouteInfo,
-	stack.PacketBufferList,
-	tcpip.NetworkProtocolNumber,
+	list stack.PacketBufferList,
 ) (int, tcpip.Error) {
-	panic("not implemented")
+	var i int = 0
+	for i, pkt := range list.AsSlice() {
+		vv := pkt.ToView()
+		n, err := vv.Read(e.writeBuf)
+		if err != nil {
+			log.Println(err)
+			return i-1, &tcpip.ErrAborted{}
+		}
+		_, err = e.ipv6rwc.Write(e.writeBuf[:n])
+		if err != nil {
+			log.Println(err)
+			return i-1, &tcpip.ErrAborted{}
+		}
+	}
+
+	return i, nil
 }
 
 func (e *YggdrasilNIC) WriteRawPacket(*stack.PacketBuffer) tcpip.Error {
@@ -130,7 +122,11 @@ func (*YggdrasilNIC) ARPHardwareType() header.ARPHardwareType {
 	return header.ARPHardwareNone
 }
 
-func (e *YggdrasilNIC) AddHeader(tcpip.LinkAddress, tcpip.LinkAddress, tcpip.NetworkProtocolNumber, *stack.PacketBuffer) {
+func (e *YggdrasilNIC) AddHeader(*stack.PacketBuffer) {
+}
+
+func (e *YggdrasilNIC) ParseHeader(*stack.PacketBuffer) bool {
+	return true
 }
 
 func (e *YggdrasilNIC) Close() error {

src/setup/setup.go

@@ -14,6 +14,9 @@ import (
 	"strings"
 
 	"github.com/gologme/log"
+	"github.com/things-go/go-socks5"
+
+	"github.com/yggdrasil-network/yggdrasil-go/contrib/netstack"
 	"github.com/yggdrasil-network/yggdrasil-go/src/address"
 	"github.com/yggdrasil-network/yggdrasil-go/src/admin"
 	"github.com/yggdrasil-network/yggdrasil-go/src/config"
@@ -21,6 +24,7 @@ import (
 	"github.com/yggdrasil-network/yggdrasil-go/src/ipv6rwc"
 	"github.com/yggdrasil-network/yggdrasil-go/src/multicast"
 	"github.com/yggdrasil-network/yggdrasil-go/src/tun"
+	"github.com/yggdrasil-network/yggdrasil-go/src/types"
 	"golang.org/x/text/encoding/unicode"
 )
 
@@ -138,9 +142,11 @@ func (n *Node) Run(args Arguments) error {
 			return err
 		}
 		address, subnet := n.core.Address(), n.core.Subnet()
-		n.logger.Infof("Your public key is %s", hex.EncodeToString(n.core.PublicKey()))
+		publicstr := hex.EncodeToString(n.core.PublicKey())
+		n.logger.Infof("Your public key is %s", publicstr)
 		n.logger.Infof("Your IPv6 address is %s", address.String())
 		n.logger.Infof("Your IPv6 subnet is %s", subnet.String())
+		n.logger.Infof("Your Yggstack resolver name is %s%s", publicstr, types.NameMappingSuffix)
 	}
 
 	// Setup the admin socket.
@@ -228,6 +234,51 @@ func (n *Node) SetupTun() error {
 	return nil
 }
 
+func (n *Node) SetupNetstack(socks *string, nameserver *string, expose *types.TCPMappings) error {
+	s, err := netstack.CreateYggdrasilNetstack(n.core)
+	if err != nil {
+		return err
+	}
+
+	// Create SOCKS server
+	if socks != nil && nameserver != nil && *socks != "" {
+		resolver := types.NewNameResolver(s, *nameserver)
+		server := socks5.NewServer(
+			socks5.WithDial(s.DialContext),
+			socks5.WithResolver(resolver),
+		)
+		go server.ListenAndServe("tcp", *socks) // nolint:errcheck
+	}
+
+	// Create TCP mappings
+	if expose != nil {
+		for _, mapping := range *expose {
+			go func(mapping types.TCPMapping) {
+				listener, err := s.ListenTCP(mapping.Listen)
+				if err != nil {
+					panic(err)
+				}
+				n.logger.Infof("Mapping Yggdrasil port %d to %s", mapping.Listen.Port, mapping.Mapped)
+				for {
+					c, err := listener.Accept()
+					if err != nil {
+						panic(err)
+					}
+					r, err := net.DialTCP("tcp", nil, mapping.Mapped)
+					if err != nil {
+						n.logger.Errorf("Failed to connect to %s: %s", mapping.Mapped, err)
+						_ = c.Close()
+						continue
+					}
+					types.ProxyTCP(n.core.MTU(), c, r)
+				}
+			}(mapping)
+		}
+	}
+
+	return nil
+}
+
 func ReadConfig(log *log.Logger, useconf bool, useconffile string, normaliseconf bool) *config.NodeConfig {
 	// Use a configuration file. If -useconf, the configuration will be read
 	// from stdin. If -useconffile, the configuration will be read from the