Merge pull request #29 from jmgraeffe/feature/cancellable-dial-in

V1: made QUIC socket dial in functions cancellable

Author: jmgraeffe

api/smp.go

@@ -1,6 +1,7 @@
 package smp
 
 import (
+	"context"
 	"time"
 
 	"github.com/netsys-lab/scion-path-discovery/packets"
@@ -476,3 +477,11 @@ func (l *MPListener) Listen() error {
 func (l *MPListener) WaitForMPPeerSockConnect() (*snet.UDPAddr, error) {
 	return l.socket.WaitForDialIn()
 }
+
+// Waits for new incoming MPPeerSocks
+// Should be called in a loop
+// Using the returned addr, a new MPPeerSock can be instantiated
+// That dials back to the incoming socket
+func (l *MPListener) WaitForMPPeerSockConnectWithContext(ctx context.Context) (*snet.UDPAddr, error) {
+	return l.socket.WaitForDialInWithContext(ctx)
+}

packets/quicconn.go

@@ -244,6 +244,24 @@ func (qc *QUICReliableConn) AcceptStream() (quic.Stream, error) {
 	return stream, nil
 }
 
+func (qc *QUICReliableConn) AcceptStreamWithContext(ctx context.Context) (quic.Stream, error) {
+	log.Debugf("Accepting on quic %s", qc.listener.Addr())
+	session, err := qc.listener.Accept(ctx)
+	if err != nil {
+		return nil, err
+	}
+	log.Debugf("Got session on quic %s", qc.listener.Addr())
+
+	stream, err := session.AcceptStream(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	// qc.internalConn = stream
+
+	return stream, nil
+}
+
 func (qc *QUICReliableConn) Listen(addr snet.UDPAddr) error {
 	qc.Ready = make(chan bool, 0)
 	udpAddr := net.UDPAddr{

socket/quicsocket.go

@@ -2,6 +2,7 @@ package socket
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"encoding/base32"
 	"encoding/gob"
@@ -178,6 +179,48 @@ func (s *QUICSocket) WaitForDialIn() (*snet.UDPAddr, error) {
 	return &addr, nil
 }
 
+func (s *QUICSocket) WaitForDialInWithContext(ctx context.Context) (*snet.UDPAddr, error) {
+	bts := make([]byte, packets.PACKET_SIZE)
+	log.Debugf("Wait for Dial In")
+	stream, err := s.listenConns[0].AcceptStreamWithContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	s.listenConns[0].SetStream(stream)
+
+	select {
+	case s.listenConns[0].Ready <- true:
+	default:
+	}
+
+	_, err = stream.Read(bts)
+	if err != nil {
+		return nil, err
+	}
+	p := DialPacketQuic{}
+	network := bytes.NewBuffer(bts) // Stand-in for a network connection
+	dec := gob.NewDecoder(network)
+	err = dec.Decode(&p)
+	if err != nil {
+		return nil, err
+	}
+
+	s.listenConns[0].SetRemote(&p.Addr)
+	log.Debugf("Waiting for %d more connections", p.NumPaths-1)
+
+	for i := 1; i < p.NumPaths; i++ {
+		_, err := s.WaitForIncomingConn()
+		if err != nil {
+			return nil, err
+		}
+		log.Debugf("Dialed In %d of %d", i, p.NumPaths)
+	}
+
+	addr := p.Addr
+	return &addr, nil
+}
+
 func (s *QUICSocket) Dial(remote snet.UDPAddr, path snet.Path, options DialOptions, i int) (packets.UDPConn, error) {
 	// appnet.SetPath(&remote, path)
 	if s.options == nil || !s.options.MultiportMode {

socket/quicsocket_test.go

@@ -1,7 +1,10 @@
 package socket
 
 import (
+	"context"
+	"errors"
 	"testing"
+	"time"
 
 	lookup "github.com/netsys-lab/scion-path-discovery/pathlookup"
 	"github.com/netsys-lab/scion-path-discovery/pathselection"
@@ -63,4 +66,22 @@ func Test_QUICSocket(t *testing.T) {
 		sock2.CloseAll()
 	})
 
+	t.Run("QUICSocket Listen And Contextualized Wait For Dial", func(t *testing.T) {
+		sock := NewQUICSocket("1-ff00:0:110,[127.0.0.12]:41000", &SockOptions{PathSelectionResponsibility: "server"})
+		err := sock.Listen()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer sock.CloseAll()
+
+		ctx, cancelFunc := context.WithDeadline(context.Background(), time.Now())
+		defer cancelFunc()
+		_, err = sock.WaitForDialInWithContext(ctx)
+		if err != nil && !errors.Is(err, context.DeadlineExceeded) {
+			t.Error(err)
+			return
+		}
+	})
+
 }

socket/scionsocket.go

@@ -2,7 +2,10 @@ package socket
 
 import (
 	"bytes"
+	"context"
 	"encoding/gob"
+	"os"
+	"time"
 
 	"github.com/netsys-lab/scion-path-discovery/packets"
 	"github.com/netsys-lab/scion-path-discovery/pathselection"
@@ -74,6 +77,34 @@ func (s *SCIONSocket) WaitForDialIn() (*snet.UDPAddr, error) {
 	return &addr, nil
 }
 
+func (s *SCIONSocket) WaitForDialInWithContext(ctx context.Context) (*snet.UDPAddr, error) {
+	for {
+		deadline, hasDeadline := ctx.Deadline()
+		var err error
+		if hasDeadline {
+			err = s.connections[0].SetReadDeadline(deadline)
+		} else {
+			err = s.connections[0].SetDeadline(time.Now().Add(5 * time.Second))
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		addr, err := s.WaitForDialIn()
+		if err == nil {
+			return addr, nil
+		} else if !os.IsTimeout(err) {
+			return nil, err
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+
+}
+
 func (s *SCIONSocket) Dial(remote snet.UDPAddr, path snet.Path, options DialOptions, i int) (packets.UDPConn, error) {
 	// appnet.SetPath(&remote, path)
 	// fmt.Printf("Dialing to %s via %s\n", remote.String(), remote.Path)

socket/scionsocket_test.go

@@ -1,7 +1,10 @@
 package socket
 
 import (
+	"context"
+	"errors"
 	"testing"
+	"time"
 
 	lookup "github.com/netsys-lab/scion-path-discovery/pathlookup"
 	"github.com/netsys-lab/scion-path-discovery/pathselection"
@@ -64,4 +67,22 @@ func Test_SCIONSocket(t *testing.T) {
 		sock2.CloseAll()
 	})
 
+	t.Run("SCIONSocket Listen And Contextualized Wait For Dial", func(t *testing.T) {
+		sock := NewSCIONSocket("1-ff00:0:110,[127.0.0.12]:21000")
+		err := sock.Listen()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer sock.CloseAll()
+
+		ctx, cancelFunc := context.WithDeadline(context.Background(), time.Now())
+		defer cancelFunc()
+		_, err = sock.WaitForDialInWithContext(ctx)
+		if err != nil && !errors.Is(err, context.DeadlineExceeded) {
+			t.Error(err)
+			return
+		}
+	})
+
 }

socket/socket.go

@@ -1,6 +1,7 @@
 package socket
 
 import (
+	"context"
 	"github.com/netsys-lab/scion-path-discovery/packets"
 	"github.com/netsys-lab/scion-path-discovery/pathselection"
 	"github.com/scionproto/scion/go/lib/snet"
@@ -26,6 +27,7 @@ type DialOptions struct {
 type UnderlaySocket interface {
 	Listen() error
 	WaitForDialIn() (*snet.UDPAddr, error)
+	WaitForDialInWithContext(ctx context.Context) (*snet.UDPAddr, error)
 	WaitForIncomingConn() (packets.UDPConn, error)
 	Dial(remote snet.UDPAddr, path snet.Path, options DialOptions, i int) (packets.UDPConn, error)
 	DialAll(remote snet.UDPAddr, path []pathselection.PathQuality, options DialOptions) ([]packets.UDPConn, error)