plugin: new framework to implement plugins

cmd/age/age_test.go

@@ -5,11 +5,11 @@
 package main
 
 import (
-	"bufio"
 	"os"
 	"testing"
 
 	"filippo.io/age"
+	"filippo.io/age/plugin"
 	"github.com/rogpeppe/go-internal/testscript"
 )
 
@@ -30,51 +30,31 @@ func TestMain(m *testing.M) {
 			return 0
 		},
 		"age-plugin-test": func() (exitCode int) {
-			// TODO: use plugin server package once it's available.
-			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")
-				return 0
-			case "--age-plugin=identity-v1":
-				scanner := bufio.NewScanner(os.Stdin)
-				scanner.Scan() // add-identity
-				scanner.Scan() // body
-				scanner.Scan() // grease
-				scanner.Scan() // body
-				scanner.Scan() // recipient-stanza
-				scanner.Scan() // body
-				fileKey := scanner.Text()
-				scanner.Scan() // done
-				scanner.Scan() // body
-				os.Stdout.WriteString("-> file-key 0\n")
-				os.Stdout.WriteString(fileKey + "\n")
-				scanner.Scan() // ok
-				scanner.Scan() // body
-				os.Stdout.WriteString("-> done\n\n")
-				return 0
-			default:
-				return 1
-			}
+			p, _ := plugin.New("test")
+			p.HandleRecipient(func(data []byte) (age.Recipient, error) {
+				return testPlugin{}, nil
+			})
+			p.HandleIdentity(func(data []byte) (age.Identity, error) {
+				return testPlugin{}, nil
+			})
+			return p.Main()
 		},
 	}))
 }
 
+type testPlugin struct{}
+
+func (testPlugin) Wrap(fileKey []byte) ([]*age.Stanza, error) {
+	return []*age.Stanza{{Type: "test", Body: fileKey}}, nil
+}
+
+func (testPlugin) Unwrap(ss []*age.Stanza) ([]byte, error) {
+	if len(ss) == 1 && ss[0].Type == "test" {
+		return ss[0].Body, nil
+	}
+	return nil, age.ErrIncorrectIdentity
+}
+
 func TestScript(t *testing.T) {
 	testscript.Run(t, testscript.Params{
 		Dir: "testdata",

plugin/client.go

@@ -79,7 +79,7 @@ func (r *Recipient) WrapWithLabels(fileKey []byte) (stanzas []*age.Stanza, label
 	if err := writeStanza(conn, addType, r.encoding); err != nil {
 		return nil, nil, err
 	}
-	if err := writeStanza(conn, fmt.Sprintf("grease-%x", rand.Int())); err != nil {
+	if _, err := writeGrease(conn); err != nil {
 		return nil, nil, err
 	}
 	if err := writeStanzaWithBody(conn, "wrap-file-key", fileKey); err != nil {
@@ -220,7 +220,7 @@ func (i *Identity) Unwrap(stanzas []*age.Stanza) (fileKey []byte, err error) {
 	if err := writeStanza(conn, "add-identity", i.encoding); err != nil {
 		return nil, err
 	}
-	if err := writeStanza(conn, fmt.Sprintf("grease-%x", rand.Int())); err != nil {
+	if _, err := writeGrease(conn); err != nil {
 		return nil, err
 	}
 	for _, rs := range stanzas {
@@ -449,3 +449,18 @@ func writeStanzaWithBody(conn io.Writer, t string, body []byte) error {
 	s := &format.Stanza{Type: t, Body: body}
 	return s.Marshal(conn)
 }
+
+func writeGrease(conn io.Writer) (sent bool, err error) {
+	if rand.Intn(3) == 0 {
+		return false, nil
+	}
+	s := &format.Stanza{Type: fmt.Sprintf("grease-%x", rand.Int())}
+	for i := 0; i < rand.Intn(3); i++ {
+		s.Args = append(s.Args, fmt.Sprintf("%d", rand.Intn(100)))
+	}
+	if rand.Intn(2) == 0 {
+		s.Body = make([]byte, rand.Intn(100))
+		rand.Read(s.Body)
+	}
+	return true, s.Marshal(conn)
+}

plugin/client_test.go

@@ -7,7 +7,6 @@
 package plugin
 
 import (
-	"bufio"
 	"io"
 	"os"
 	"path/filepath"
@@ -20,63 +19,41 @@ import (
 
 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])
-		}
+		p, _ := New("test")
+		p.HandleRecipient(func(data []byte) (age.Recipient, error) {
+			return testRecipient{}, nil
+		})
+		os.Exit(p.Main())
 	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])
-		}
+		p, _ := New("testpqc")
+		p.HandleRecipient(func(data []byte) (age.Recipient, error) {
+			return testPQCRecipient{}, nil
+		})
+		os.Exit(p.Main())
 	default:
 		os.Exit(m.Run())
 	}
 }
 
+type testRecipient struct{}
+
+func (testRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
+	return []*age.Stanza{{Type: "test", Body: fileKey}}, nil
+}
+
+type testPQCRecipient struct{}
+
+var _ age.RecipientWithLabels = testPQCRecipient{}
+
+func (testPQCRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
+	return []*age.Stanza{{Type: "test", Body: fileKey}}, nil
+}
+
+func (testPQCRecipient) WrapWithLabels(fileKey []byte) ([]*age.Stanza, []string, error) {
+	return []*age.Stanza{{Type: "test", Body: fileKey}}, []string{"postquantum"}, nil
+}
+
 func TestLabels(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("Windows support is TODO")

plugin/example_test.go

@@ -0,0 +1,43 @@
+package plugin_test
+
+import (
+	"log"
+	"os"
+
+	"filippo.io/age"
+	"filippo.io/age/plugin"
+)
+
+type Recipient struct{}
+
+func (r *Recipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
+	panic("unimplemented")
+}
+
+func NewRecipient(data []byte) (*Recipient, error) {
+	return &Recipient{}, nil
+}
+
+type Identity struct{}
+
+func (i *Identity) Unwrap(s []*age.Stanza) ([]byte, error) {
+	panic("unimplemented")
+}
+
+func NewIdentity(data []byte) (*Identity, error) {
+	return &Identity{}, nil
+}
+
+func ExamplePlugin_main() {
+	p, err := plugin.New("example")
+	if err != nil {
+		log.Fatal(err)
+	}
+	p.HandleRecipient(func(data []byte) (age.Recipient, error) {
+		return NewRecipient(data)
+	})
+	p.HandleIdentity(func(data []byte) (age.Identity, error) {
+		return NewIdentity(data)
+	})
+	os.Exit(p.Main())
+}