two-memnet-listeners.go

package main

import (
	"cmp"
	"context"
	"fmt"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"time"

	"tailscale.com/net/memnet"
)

type tunnel struct {
	net      *memnet.Network
	listener *listener
}

func (l *tunnel) StartListener(endpoint netip.AddrPort) (net.Listener, error) {
	l.listener = newListener(l.net)
	return l.listener, l.listener.StartExternal(endpoint)
}

func (l *tunnel) AddInternalIPToListener(endpoint netip.AddrPort) error {
	if l.listener == nil {
		return os.ErrNotExist

	}
	return l.listener.StartInternal(endpoint)
}

// newListener creates a new listener instance handling two endpoints in a single listener.
func newListener(mnet *memnet.Network) *listener {
	return &listener{
		net:     mnet,
		connCh:  make(chan net.Conn),
		errCh:   make(chan error, 1),
		closeCh: make(chan struct{}),
	}
}

var _ net.Listener = (*listener)(nil)

// listener handles listening on two endpoints in a single listener.
type listener struct {
	net                *memnet.Network
	internal, external net.Listener

	connCh  chan net.Conn
	errCh   chan error
	closeCh chan struct{}
}

// StartExternal starts listening on the external endpoint.
func (l *listener) StartExternal(endpoint netip.AddrPort) (err error) {
	l.external, err = l.net.Listen("tcp", endpoint.String())
	if err != nil {
		slog.Error("Failed to start external memory listener", slog.String("endpoint", endpoint.String()))
		return err
	}

	go l.serve(l.external)
	return nil
}

// StartInternal starts listening on the internal endpoint.
func (l *listener) StartInternal(endpoint netip.AddrPort) (err error) {
	l.internal, err = l.net.Listen("tcp", endpoint.String())
	if err != nil {
		slog.Error("Failed to start internal memory listener", slog.String("endpoint", endpoint.String()))
		return err
	}

	go l.serve(l.internal)
	return nil
}

// serve runs a listener and forwards incoming connections.
func (l *listener) serve(ln net.Listener) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			l.errCh <- err
			return
		}

		select {
		case l.connCh <- conn:
		case <-l.closeCh:
			return
		}
	}
}

// Accept accepts incoming connections and errors from the listeners and handles closing as well.
func (l *listener) Accept() (net.Conn, error) {
	select {
	case conn := <-l.connCh:
		return conn, nil
	case err := <-l.errCh:
		return nil, err
	case <-l.closeCh:
		return nil, nil
	}
}

// Close closes the listener and returns any errors encountered.
func (l *listener) Close() error {
	close(l.closeCh)
	var erri, erre error
	if l.internal != nil {
		erri = l.internal.Close()
	}
	if l.external != nil {
		erre = l.external.Close()
	}
	return cmp.Or(erri, erre)
}

// Addr returns the address of the listener.
func (l *listener) Addr() net.Addr {
	return l.external.Addr()
}

func main() {
	t := &tunnel{
		net: &memnet.Network{},
	}
	listener, err := t.StartListener(netip.MustParseAddrPort("10.0.0.1:50000"))
	if err != nil {
		panic(err)
	}

	go func() {
		time.Sleep(1 * time.Second)
		conn, err := t.net.Dial(context.Background(), "tcp", "10.0.0.1:50000")
		if err != nil {
			fmt.Println("Error dialing 1:", err)
			return
		}
		fmt.Println("Dial 1 successful:", conn.Close())

		time.Sleep(2 * time.Second)
		err = t.AddInternalIPToListener(netip.MustParseAddrPort("100.0.0.1:50000"))
		if err != nil {
			fmt.Println("Error adding internal IP:", err)
			return
		}

		time.Sleep(1 * time.Second)

		conn, err = t.net.Dial(context.Background(), "tcp", "100.0.0.1:50000")
		if err != nil {
			fmt.Println("Error dialing 2:", err)
			return
		}
		fmt.Println("Dial 2 successful:", conn.Close())

		listener.Close()
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			break
		} else if conn == nil {
			break
		}
		go handleConnection(conn)
	}

	fmt.Println("Clean exit")
}

func handleConnection(conn net.Conn) {
	defer conn.Close()
	fmt.Println("Accepted connection from", conn.RemoteAddr())
}