feat: add support for libp2p pnet

Author: aschmahmann

lib/bitswap.go

@@ -61,7 +61,7 @@ var _ json.Marshaler = (*BsCheckOutput)(nil)
 func CheckBitswapCID(ctx context.Context, h host.Host, c cid.Cid, ma multiaddr.Multiaddr, getBlock bool) (*BsCheckOutput, error) {
 	var err error
 	if h == nil {
-		h, err = libp2pHost()
+		h, err = libp2pHost(ctx)
 	}
 
 	if err != nil {
@@ -171,10 +171,8 @@ loop:
 	}, nil
 }
 
-func GetBitswapCID(root cid.Cid, ai *peer.AddrInfo) error {
-
-	ctx := context.Background()
-	h, err := libp2pHost()
+func GetBitswapCID(ctx context.Context, root cid.Cid, ai *peer.AddrInfo) error {
+	h, err := libp2pHost(ctx)
 	if err != nil {
 		return err
 	}

lib/dht.go

@@ -16,7 +16,7 @@ import (
 )
 
 func DhtProtocolMessenger(ctx context.Context, proto protocol.ID, ai *peer.AddrInfo) (*dhtpb.ProtocolMessenger, error) {
-	h, err := libp2pHost()
+	h, err := libp2pHost(ctx)
 	if err != nil {
 		return nil, err
 	}

lib/identify.go

@@ -45,7 +45,7 @@ func IdentifyRequest(ctx context.Context, maStr string, allowUnknownPeer bool) (
 		}
 	}
 
-	h, err := libp2pHost()
+	h, err := libp2pHost(ctx)
 	if err != nil {
 		return nil, err
 	}

lib/libp2p.go

@@ -1,14 +1,21 @@
 package vole
 
 import (
+	"context"
+
 	"github.com/libp2p/go-libp2p"
 	"github.com/libp2p/go-libp2p/core/host"
 )
 
-func libp2pHost() (host.Host, error) {
-	h, err := libp2p.New(
+func libp2pHost(ctx context.Context) (host.Host, error) {
+	opts := []libp2p.Option{
 		libp2p.EnableHolePunching(),
-	)
+	}
+	if psk, ok := pnetPSKFromContext(ctx); ok {
+		opts = append(opts, libp2p.PrivateNetwork(psk))
+	}
+
+	h, err := libp2p.New(opts...)
 	if err != nil {
 		return nil, err
 	}

lib/only_connect.go

@@ -11,7 +11,7 @@ import (
 )
 
 func OnlyConnect(ctx context.Context, p *peer.AddrInfo) error {
-	h, err := libp2pHost()
+	h, err := libp2pHost(ctx)
 	if err != nil {
 		return err
 	}

lib/ping.go

@@ -26,7 +26,7 @@ func Ping(ctx context.Context, forceRelay bool, p *peer.AddrInfo) error {
 		}
 	}
 
-	h, err := libp2pHost()
+	h, err := libp2pHost(ctx)
 	if err != nil {
 		return err
 	}

lib/pnet.go

@@ -0,0 +1,42 @@
+package vole
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/libp2p/go-libp2p/core/pnet"
+)
+
+type pnetPSKContextKey struct{}
+
+func WithPnetPSK(ctx context.Context, psk pnet.PSK) context.Context {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	cpy := make([]byte, len(psk))
+	copy(cpy, psk)
+	return context.WithValue(ctx, pnetPSKContextKey{}, pnet.PSK(cpy))
+}
+
+func pnetPSKFromContext(ctx context.Context) (pnet.PSK, bool) {
+	if ctx == nil {
+		return nil, false
+	}
+	psk, ok := ctx.Value(pnetPSKContextKey{}).(pnet.PSK)
+	return psk, ok
+}
+
+func LoadPnetPSK(path string) (pnet.PSK, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("read pnet swarm key %q: %w", path, err)
+	}
+	defer func() { _ = f.Close() }()
+
+	psk, err := pnet.DecodeV1PSK(f)
+	if err != nil {
+		return nil, fmt.Errorf("decode pnet swarm key %q: %w", path, err)
+	}
+	return psk, nil
+}

lib/pnet_test.go

@@ -0,0 +1,141 @@
+package vole
+
+import (
+	"context"
+	"errors"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/libp2p/go-libp2p/core/peer"
+	corepnet "github.com/libp2p/go-libp2p/core/pnet"
+)
+
+const testSwarmKeyV1 = "/key/swarm/psk/1.0.0/\n/base16/\n0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef\n"
+const testSwarmKeyV1Alt = "/key/swarm/psk/1.0.0/\n/base16/\nfedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210\n"
+
+func writeTempFile(t *testing.T, contents string) string {
+	t.Helper()
+	f, err := os.CreateTemp(t.TempDir(), "swarmkey-*")
+	if err != nil {
+		t.Fatalf("CreateTemp: %v", err)
+	}
+	if _, err := f.WriteString(contents); err != nil {
+		_ = f.Close()
+		t.Fatalf("WriteString: %v", err)
+	}
+	if err := f.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	return f.Name()
+}
+
+func TestLoadPnetPSK_MissingFile(t *testing.T) {
+	_, err := LoadPnetPSK("does-not-exist")
+	if err == nil {
+		t.Fatalf("expected error")
+	}
+}
+
+func TestLoadPnetPSK_InvalidFile(t *testing.T) {
+	path := writeTempFile(t, "not a swarm key")
+	_, err := LoadPnetPSK(path)
+	if err == nil {
+		t.Fatalf("expected error")
+	}
+}
+
+func TestPnet_AllowsOnlySameKey(t *testing.T) {
+	path := writeTempFile(t, testSwarmKeyV1)
+	psk, err := LoadPnetPSK(path)
+	if err != nil {
+		t.Fatalf("LoadPnetPSK: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	hNo, err := libp2pHost(ctx)
+	if err != nil {
+		t.Fatalf("libp2pHost(no pnet): %v", err)
+	}
+	defer hNo.Close()
+
+	hYesA, err := libp2pHost(WithPnetPSK(ctx, psk))
+	if err != nil {
+		t.Fatalf("libp2pHost(pnet A): %v", err)
+	}
+	defer hYesA.Close()
+
+	hYesB, err := libp2pHost(WithPnetPSK(ctx, psk))
+	if err != nil {
+		t.Fatalf("libp2pHost(pnet B): %v", err)
+	}
+	defer hYesB.Close()
+
+	// pnet->no-pnet should fail.
+	aiNo := peer.AddrInfo{ID: hNo.ID(), Addrs: hNo.Addrs()}
+	dialCtx, dialCancel := context.WithTimeout(ctx, 3*time.Second)
+	err = hYesA.Connect(dialCtx, aiNo)
+	dialCancel()
+	if err == nil {
+		t.Fatalf("expected pnet dial to non-pnet host to fail")
+	}
+	// Depending on the transport/OS, the pnet failure may be wrapped in dial/security negotiation errors.
+	if !corepnet.IsPNetError(err) && !strings.Contains(strings.ToLower(err.Error()), "privnet") {
+		// Don't fail the test on classification; the key property is that the dial fails.
+		t.Logf("dial failed with non-pnet-classified error: %T: %v", err, err)
+	}
+
+	// pnet->pnet (same key) should succeed.
+	aiYesB := peer.AddrInfo{ID: hYesB.ID(), Addrs: hYesB.Addrs()}
+	dialCtx2, dialCancel2 := context.WithTimeout(ctx, 5*time.Second)
+	err = hYesA.Connect(dialCtx2, aiYesB)
+	dialCancel2()
+	if err != nil {
+		// If this fails, surface the most useful root cause.
+		var pnetErr corepnet.Error
+		if errors.As(err, &pnetErr) {
+			t.Fatalf("unexpected pnet error connecting same-key hosts: %v", err)
+		}
+		t.Fatalf("expected same-key hosts to connect, got: %v", err)
+	}
+}
+
+func TestPnet_DifferentKeysFail(t *testing.T) {
+	pathA := writeTempFile(t, testSwarmKeyV1)
+	pskA, err := LoadPnetPSK(pathA)
+	if err != nil {
+		t.Fatalf("LoadPnetPSK(A): %v", err)
+	}
+
+	pathB := writeTempFile(t, testSwarmKeyV1Alt)
+	pskB, err := LoadPnetPSK(pathB)
+	if err != nil {
+		t.Fatalf("LoadPnetPSK(B): %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	hA, err := libp2pHost(WithPnetPSK(ctx, pskA))
+	if err != nil {
+		t.Fatalf("libp2pHost(pnet A): %v", err)
+	}
+	defer hA.Close()
+
+	hB, err := libp2pHost(WithPnetPSK(ctx, pskB))
+	if err != nil {
+		t.Fatalf("libp2pHost(pnet B): %v", err)
+	}
+	defer hB.Close()
+
+	aiB := peer.AddrInfo{ID: hB.ID(), Addrs: hB.Addrs()}
+	dialCtx, dialCancel := context.WithTimeout(ctx, 3*time.Second)
+	err = hA.Connect(dialCtx, aiB)
+	dialCancel()
+	if err == nil {
+		t.Fatalf("expected different-key pnet connect to fail")
+	}
+}

main.go

@@ -37,12 +37,31 @@ func main() {
 	app := &cli.App{
 		Name:  "vole",
 		Usage: "a collection of tools for digging around IPFS nodes",
+		Before: func(c *cli.Context) error {
+			pnetPath := c.Path("pnet")
+			if pnetPath == "" {
+				return nil
+			}
+			psk, err := vole.LoadPnetPSK(pnetPath)
+			if err != nil {
+				return err
+			}
+			c.Context = vole.WithPnetPSK(c.Context, psk)
+			return nil
+		},
 		Authors: []*cli.Author{
 			{
 				Name:  "Adin Schmahmann",
 				Email: "[email protected]",
 			},
 		},
+		Flags: []cli.Flag{
+			&cli.PathFlag{
+				Name:  "pnet",
+				Usage: "path to a libp2p private network swarm key file to use for all commands that create a libp2p host",
+				Value: "",
+			},
+		},
 		Commands: []*cli.Command{
 			{
 				Name:  "bitswap",
@@ -354,7 +373,6 @@ Note: may not work with some transports such as p2p-circuit (not applicable) and
 						Name:      "ping",
 						ArgsUsage: "<multiaddr>",
 						Flags: []cli.Flag{
-
 							&cli.BoolFlag{
 								Name:        "force-relay",
 								Usage:       `Ping the peer over a relay instead of a direct connection`,
@@ -446,7 +464,7 @@ var bitswapGetCmd = &cli.Command{
 			return err
 		}
 
-		return vole.GetBitswapCID(root, ai)
+		return vole.GetBitswapCID(cctx.Context, root, ai)
 	},
 }
 var bitswapCheckCmd = &cli.Command{