feat(sdk): stub ipc

Author: joe0BAB

pkg/stub/ipc.go

@@ -0,0 +1,104 @@
+package stub
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"sync"
+
+	"github.com/containerd/nri/pkg/net/multiplex"
+)
+
+type ipcImpl struct {
+	mConn     net.Conn
+	doneC     chan struct{}
+	serverErr error
+	teardown  func() error
+}
+
+type Ipc interface {
+	conn() net.Conn
+	wait(ctx context.Context) error
+	close() error
+}
+
+func NewIPC(sockConn net.Conn, handler http.Handler) (result Ipc, retErr error) {
+	mux := multiplex.Multiplex(sockConn)
+	defer func() {
+		if retErr != nil {
+			mux.Close()
+		}
+	}()
+	doneC := make(chan struct{})
+	listener, err := mux.Listen(multiplex.PluginServiceConn)
+	if err != nil {
+		listener.Close()
+		return nil, err
+	}
+	server := &http.Server{
+		Handler: handler,
+	}
+	var serverErr error
+	go func() {
+		err := server.Serve(listener)
+		if !errors.Is(err, http.ErrServerClosed) {
+			serverErr = err
+		}
+		close(doneC)
+	}()
+	defer func() {
+		if retErr != nil {
+			server.Close()
+		}
+	}()
+	conn, err := mux.Open(multiplex.RuntimeServiceConn)
+	if err != nil {
+		return nil, fmt.Errorf("failed to multiplex grcp client connection: %w", err)
+	}
+	defer func() {
+		if retErr != nil {
+			_ = conn.Close()
+		}
+	}()
+	return &ipcImpl{
+		mConn:     conn,
+		doneC:     doneC,
+		serverErr: serverErr,
+		teardown: sync.OnceValue(func() error {
+			var errList []error
+			if err := listener.Close(); err != nil {
+				errList = append(errList, err)
+			}
+			if err := server.Close(); err != nil {
+				errList = append(errList, err)
+			}
+			if err := conn.Close(); err != nil {
+				errList = append(errList, err)
+			}
+			if err := mux.Close(); err != nil {
+				errList = append(errList, err)
+			}
+			<-doneC
+			return errors.Join(errList...)
+		}),
+	}, nil
+}
+
+func (i *ipcImpl) conn() net.Conn {
+	return i.mConn
+}
+
+func (i *ipcImpl) wait(ctx context.Context) error {
+	select {
+	case <-i.doneC:
+		return i.serverErr
+	case <-ctx.Done():
+		return nil
+	}
+}
+
+func (i *ipcImpl) close() error {
+	return i.teardown()
+}