age,plugin: add RecipientWithLabels

plugin/client.go

@@ -14,6 +14,7 @@ import (
 	"io"
 	"math/rand"
 	"os"
+	"path/filepath"
 	"strconv"
 	"time"
 
@@ -33,6 +34,7 @@ type Recipient struct {
 }
 
 var _ age.Recipient = &Recipient{}
+var _ age.RecipientWithLabels = &Recipient{}
 
 func NewRecipient(s string, ui *ClientUI) (*Recipient, error) {
 	name, _, err := ParseRecipient(s)
@@ -52,6 +54,11 @@ func (r *Recipient) Name() string {
 }
 
 func (r *Recipient) Wrap(fileKey []byte) (stanzas []*age.Stanza, err error) {
+	stanzas, _, err = r.WrapWithLabels(fileKey)
+	return
+}
+
+func (r *Recipient) WrapWithLabels(fileKey []byte) (stanzas []*age.Stanza, labels []string, err error) {
 	defer func() {
 		if err != nil {
 			err = fmt.Errorf("%s plugin: %w", r.name, err)
@@ -60,7 +67,7 @@ func (r *Recipient) Wrap(fileKey []byte) (stanzas []*age.Stanza, err error) {
 
 	conn, err := openClientConnection(r.name, "recipient-v1")
 	if err != nil {
-		return nil, fmt.Errorf("couldn't start plugin: %v", err)
+		return nil, nil, fmt.Errorf("couldn't start plugin: %v", err)
 	}
 	defer conn.Close()
 
@@ -70,16 +77,19 @@ func (r *Recipient) Wrap(fileKey []byte) (stanzas []*age.Stanza, err error) {
 		addType = "add-identity"
 	}
 	if err := writeStanza(conn, addType, r.encoding); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	if err := writeStanza(conn, fmt.Sprintf("grease-%x", rand.Int())); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	if err := writeStanzaWithBody(conn, "wrap-file-key", fileKey); err != nil {
-		return nil, err
+		return nil, nil, err
+	}
+	if err := writeStanza(conn, "extension-labels"); err != nil {
+		return nil, nil, err
 	}
 	if err := writeStanza(conn, "done"); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// Phase 2: plugin responds with stanzas
@@ -88,21 +98,21 @@ ReadLoop:
 	for {
 		s, err := r.ui.readStanza(r.name, sr)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 
 		switch s.Type {
 		case "recipient-stanza":
 			if len(s.Args) < 2 {
-				return nil, fmt.Errorf("malformed recipient stanza: unexpected argument count")
+				return nil, nil, fmt.Errorf("malformed recipient stanza: unexpected argument count")
 			}
 			n, err := strconv.Atoi(s.Args[0])
 			if err != nil {
-				return nil, fmt.Errorf("malformed recipient stanza: invalid index")
+				return nil, nil, fmt.Errorf("malformed recipient stanza: invalid index")
 			}
 			// We only send a single file key, so the index must be 0.
 			if n != 0 {
-				return nil, fmt.Errorf("malformed recipient stanza: unexpected index")
+				return nil, nil, fmt.Errorf("malformed recipient stanza: unexpected index")
 			}
 
 			stanzas = append(stanzas, &age.Stanza{
@@ -112,32 +122,41 @@ ReadLoop:
 			})
 
 			if err := writeStanza(conn, "ok"); err != nil {
-				return nil, err
+				return nil, nil, err
+			}
+		case "labels":
+			if labels != nil {
+				return nil, nil, fmt.Errorf("repeated labels stanza")
+			}
+			labels = s.Args
+
+			if err := writeStanza(conn, "ok"); err != nil {
+				return nil, nil, err
 			}
 		case "error":
 			if err := writeStanza(conn, "ok"); err != nil {
-				return nil, err
+				return nil, nil, err
 			}
 
-			return nil, fmt.Errorf("%s", s.Body)
+			return nil, nil, fmt.Errorf("%s", s.Body)
 		case "done":
 			break ReadLoop
 		default:
 			if ok, err := r.ui.handle(r.name, conn, s); err != nil {
-				return nil, err
+				return nil, nil, err
 			} else if !ok {
 				if err := writeStanza(conn, "unsupported"); err != nil {
-					return nil, err
+					return nil, nil, err
 				}
 			}
 		}
 	}
 
 	if len(stanzas) == 0 {
-		return nil, fmt.Errorf("received zero recipient stanzas")
+		return nil, nil, fmt.Errorf("received zero recipient stanzas")
 	}
 
-	return stanzas, nil
+	return stanzas, labels, nil
 }
 
 type Identity struct {
@@ -367,8 +386,14 @@ type clientConnection struct {
 	close     func()
 }
 
+var testOnlyPluginPath string
+
 func openClientConnection(name, protocol string) (*clientConnection, error) {
-	cmd := exec.Command("age-plugin-"+name, "--age-plugin="+protocol)
+	path := "age-plugin-" + name
+	if testOnlyPluginPath != "" {
+		path = filepath.Join(testOnlyPluginPath, path)
+	}
+	cmd := exec.Command(path, "--age-plugin="+protocol)
 
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {

plugin/client_test.go

@@ -0,0 +1,129 @@
+// Copyright 2023 The age Authors
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+package plugin
+
+import (
+	"bufio"
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"filippo.io/age"
+	"filippo.io/age/internal/bech32"
+)
+
+func TestMain(m *testing.M) {
+	switch filepath.Base(os.Args[0]) {
+	// TODO: deduplicate from cmd/age TestMain.
+	case "age-plugin-test":
+		switch os.Args[1] {
+		case "--age-plugin=recipient-v1":
+			scanner := bufio.NewScanner(os.Stdin)
+			scanner.Scan() // add-recipient
+			scanner.Scan() // body
+			scanner.Scan() // grease
+			scanner.Scan() // body
+			scanner.Scan() // wrap-file-key
+			scanner.Scan() // body
+			fileKey := scanner.Text()
+			scanner.Scan() // extension-labels
+			scanner.Scan() // body
+			scanner.Scan() // done
+			scanner.Scan() // body
+			os.Stdout.WriteString("-> recipient-stanza 0 test\n")
+			os.Stdout.WriteString(fileKey + "\n")
+			scanner.Scan() // ok
+			scanner.Scan() // body
+			os.Stdout.WriteString("-> done\n\n")
+			os.Exit(0)
+		default:
+			panic(os.Args[1])
+		}
+	case "age-plugin-testpqc":
+		switch os.Args[1] {
+		case "--age-plugin=recipient-v1":
+			scanner := bufio.NewScanner(os.Stdin)
+			scanner.Scan() // add-recipient
+			scanner.Scan() // body
+			scanner.Scan() // grease
+			scanner.Scan() // body
+			scanner.Scan() // wrap-file-key
+			scanner.Scan() // body
+			fileKey := scanner.Text()
+			scanner.Scan() // extension-labels
+			scanner.Scan() // body
+			scanner.Scan() // done
+			scanner.Scan() // body
+			os.Stdout.WriteString("-> recipient-stanza 0 test\n")
+			os.Stdout.WriteString(fileKey + "\n")
+			scanner.Scan() // ok
+			scanner.Scan() // body
+			os.Stdout.WriteString("-> labels postquantum\n\n")
+			scanner.Scan() // ok
+			scanner.Scan() // body
+			os.Stdout.WriteString("-> done\n\n")
+			os.Exit(0)
+		default:
+			panic(os.Args[1])
+		}
+	default:
+		os.Exit(m.Run())
+	}
+}
+
+func TestLabels(t *testing.T) {
+	temp := t.TempDir()
+	testOnlyPluginPath = temp
+	t.Cleanup(func() { testOnlyPluginPath = "" })
+	ex, err := os.Executable()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Link(ex, filepath.Join(temp, "age-plugin-test")); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Chmod(filepath.Join(temp, "age-plugin-test"), 0755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Link(ex, filepath.Join(temp, "age-plugin-testpqc")); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Chmod(filepath.Join(temp, "age-plugin-testpqc"), 0755); err != nil {
+		t.Fatal(err)
+	}
+
+	name, err := bech32.Encode("age1test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	testPlugin, err := NewRecipient(name, &ClientUI{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	namePQC, err := bech32.Encode("age1testpqc", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	testPluginPQC, err := NewRecipient(namePQC, &ClientUI{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := age.Encrypt(io.Discard, testPluginPQC); err != nil {
+		t.Errorf("expected one pqc to work, got %v", err)
+	}
+	if _, err := age.Encrypt(io.Discard, testPluginPQC, testPluginPQC); err != nil {
+		t.Errorf("expected two pqc to work, got %v", err)
+	}
+	if _, err := age.Encrypt(io.Discard, testPluginPQC, testPlugin); err == nil {
+		t.Errorf("expected one pqc and one normal to fail")
+	}
+	if _, err := age.Encrypt(io.Discard, testPlugin, testPluginPQC); err == nil {
+		t.Errorf("expected one pqc and one normal to fail")
+	}
+}