feat: adds context support

Author: gomaja

sctp.go

@@ -17,8 +17,10 @@ package sctp
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
+	"log"
 	"net"
 	"strconv"
 	"strings"
@@ -437,6 +439,7 @@ func SCTPBind(fd int, addr *SCTPAddr, flags int) error {
 type SCTPConn struct {
 	_fd                 int32
 	notificationHandler NotificationHandler
+	ctx                 context.Context
 }
 
 func (c *SCTPConn) fd() int {
@@ -451,6 +454,15 @@ func NewSCTPConn(fd int, handler NotificationHandler) *SCTPConn {
 	return conn
 }
 
+func dialSCTPConnect(sock int, raddr *SCTPAddr, notificationHandler NotificationHandler) (*SCTPConn, error) {
+	_, err := SCTPConnect(sock, raddr)
+	if err != nil {
+		return nil, err
+	}
+	log.Println("4-CANCELLEDDDDDDDD")
+	return NewSCTPConn(sock, notificationHandler), nil
+}
+
 func (c *SCTPConn) Write(b []byte) (int, error) {
 	return c.SCTPWrite(b, nil)
 }

sctp_linux.go

@@ -19,7 +19,9 @@
 package sctp
 
 import (
+	"context"
 	"io"
+	"log"
 	"net"
 	"os"
 	"runtime"
@@ -105,7 +107,24 @@ func (c *SCTPConn) SCTPWrite(b []byte, info *SndRcvInfo) (int, error) {
 		hdr.SetLen(syscall.CmsgSpace(len(cmsgBuf)))
 		cbuf = append(toBuf(hdr), cmsgBuf...)
 	}
-	return syscall.SendmsgN(c.fd(), b, cbuf, nil, syscall.MSG_DONTWAIT)
+
+	if c.ctx != nil {
+		for {
+			select {
+			case <-c.ctx.Done():
+				c.Close()
+				return 0, syscall.EBADF
+			default:
+				n, err := syscall.SendmsgN(c.fd(), b, cbuf, nil, syscall.MSG_DONTWAIT)
+				if err != nil {
+					continue
+				}
+				return n, err
+			}
+		}
+	}
+
+	return syscall.SendmsgN(c.fd(), b, cbuf, nil, 0)
 }
 
 func parseSndRcvInfo(b []byte) (*SndRcvInfo, error) {
@@ -354,21 +373,67 @@ func DialSCTP(net string, laddr, raddr *SCTPAddr) (*SCTPConn, error) {
 	return DialSCTPExt(net, laddr, raddr, InitMsg{NumOstreams: SCTP_MAX_STREAM})
 }
 
+// DialSCTPWithContext - bind socket to laddr (if given) and connect to raddr with context
+func DialSCTPWithContext(net string, laddr, raddr *SCTPAddr, ctx context.Context) (*SCTPConn, error) {
+	return dialSCTPExtConfigWithContext(net, laddr, raddr, InitMsg{NumOstreams: SCTP_MAX_STREAM}, nil, nil, ctx)
+}
+
 // DialSCTPExt - same as DialSCTP but with given SCTP options
 func DialSCTPExt(network string, laddr, raddr *SCTPAddr, options InitMsg) (*SCTPConn, error) {
 	return dialSCTPExtConfig(network, laddr, raddr, options, nil, nil)
 }
 
 // dialSCTPExtConfig - same as DialSCTP but with given SCTP options and socket configuration
 func dialSCTPExtConfig(network string, laddr, raddr *SCTPAddr, options InitMsg, control func(network, address string, c syscall.RawConn) error, notificationHandler NotificationHandler) (*SCTPConn, error) {
+	sock, err := clientSocket(network, laddr, raddr, options, control)
+	if err != nil {
+		return nil, err
+	}
+
+	return dialSCTPConnect(sock, raddr, notificationHandler)
+}
+
+// dialSCTPExtConfig - same as DialSCTP but with given SCTP options and socket configuration
+func dialSCTPExtConfigWithContext(network string, laddr, raddr *SCTPAddr, options InitMsg, control func(network, address string, c syscall.RawConn) error, notificationHandler NotificationHandler, ctx context.Context) (*SCTPConn, error) {
+	sock, err := clientSocket(network, laddr, raddr, options, control)
+	if err != nil {
+		return nil, err
+	}
+
+	connChan := make(chan *SCTPConn, 1)
+	errChan := make(chan error, 1)
+	done := make(chan struct{})
+
+	go func() {
+		connection, errDial := dialSCTPConnect(sock, raddr, notificationHandler)
+		if errDial == nil {
+			connection.ctx = ctx
+		}
+		connChan <- connection
+		errChan <- errDial
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		log.Println("2-CANCELLEDDDDDDDD")
+		closeSctpSocket(sock, 1*time.Second)
+		log.Println("3-CANCELLEDDDDDDDD")
+		return nil, ctx.Err()
+	case <-done:
+		return <-connChan, <-errChan
+	}
+}
+
+func clientSocket(network string, laddr, raddr *SCTPAddr, options InitMsg, control func(network, address string, c syscall.RawConn) error) (int, error) {
 	af, ipv6only := favoriteAddrFamily(network, laddr, raddr, "dial")
 	sock, err := syscall.Socket(
 		af,
 		syscall.SOCK_STREAM|syscall.SOCK_CLOEXEC,
 		syscall.IPPROTO_SCTP,
 	)
 	if err != nil {
-		return nil, err
+		return -1, err
 	}
 
 	// close socket on error
@@ -378,7 +443,7 @@ func dialSCTPExtConfig(network string, laddr, raddr *SCTPAddr, options InitMsg,
 		}
 	}()
 	if err = setDefaultSockopts(sock, af, ipv6only); err != nil {
-		return nil, err
+		return -1, err
 	}
 	if control != nil {
 		rc := rawConn{sockfd: sock}
@@ -387,12 +452,12 @@ func dialSCTPExtConfig(network string, laddr, raddr *SCTPAddr, options InitMsg,
 			localAddressString = laddr.String()
 		}
 		if err = control(network, localAddressString, rc); err != nil {
-			return nil, err
+			return -1, err
 		}
 	}
 	err = setInitOpts(sock, options)
 	if err != nil {
-		return nil, err
+		return -1, err
 	}
 	if laddr != nil {
 		// If IP address and/or port was not provided so far, let's use the unspecified IPv4 or IPv6 address
@@ -405,12 +470,8 @@ func dialSCTPExtConfig(network string, laddr, raddr *SCTPAddr, options InitMsg,
 		}
 		err = SCTPBind(sock, laddr, SCTP_BINDX_ADD_ADDR) // error EADDRINUSE "Address already in use" may occur if resource (source IP and Port) is occupied
 		if err != nil {
-			return nil, err
+			return -1, err
 		}
 	}
-	_, err = SCTPConnect(sock, raddr)
-	if err != nil {
-		return nil, err
-	}
-	return NewSCTPConn(sock, notificationHandler), nil
+	return sock, nil
 }