plugin: implement identities in framework

Updates #485

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_test.go

@@ -19,19 +19,18 @@ import (
 
 func TestMain(m *testing.M) {
 	switch filepath.Base(os.Args[0]) {
-	// TODO: deduplicate from cmd/age TestMain.
 	case "age-plugin-test":
 		p, _ := New("test")
 		p.HandleRecipient(func(data []byte) (age.Recipient, error) {
 			return testRecipient{}, nil
 		})
-		p.Main()
+		os.Exit(p.Main())
 	case "age-plugin-testpqc":
 		p, _ := New("testpqc")
 		p.HandleRecipient(func(data []byte) (age.Recipient, error) {
 			return testPQCRecipient{}, nil
 		})
-		p.Main()
+		os.Exit(p.Main())
 	default:
 		os.Exit(m.Run())
 	}

plugin/plugin.go

@@ -2,9 +2,11 @@ package plugin
 
 import (
 	"bufio"
+	"errors"
 	"flag"
 	"fmt"
 	"os"
+	"strconv"
 
 	"filippo.io/age"
 	"filippo.io/age/internal/format"
@@ -57,7 +59,7 @@ func (p *Plugin) HandleIdentity(f func(data []byte) (age.Identity, error)) {
 	p.identity = f
 }
 
-func (p *Plugin) Main() {
+func (p *Plugin) Main() int {
 	if p.fs == nil {
 		p.RegisterFlags(nil)
 	}
@@ -65,16 +67,17 @@ func (p *Plugin) Main() {
 		p.fs.Parse(os.Args[1:])
 	}
 	if *p.sm == "recipient-v1" {
-		p.RecipientV1()
+		return p.RecipientV1()
 	}
 	if *p.sm == "identity-v1" {
-		p.IdentityV1()
+		return p.IdentityV1()
 	}
+	return fatalf("unknown state machine %q", *p.sm)
 }
 
-func (p *Plugin) RecipientV1() {
-	if p.recipient == nil {
-		fatalf("recipient-v1 not supported")
+func (p *Plugin) RecipientV1() int {
+	if p.recipient == nil && p.idAsRecipient == nil {
+		return fatalf("recipient-v1 not supported")
 	}
 
 	var recipientStrings, identityStrings []string
@@ -86,24 +89,34 @@ ReadLoop:
 	for {
 		s, err := sr.ReadStanza()
 		if err != nil {
-			fatalf("failed to read stanza: %v", err)
+			return fatalf("failed to read stanza: %v", err)
 		}
 
 		switch s.Type {
 		case "add-recipient":
-			expectStanzaWithNoBody(s, 1)
+			if err := expectStanzaWithNoBody(s, 1); err != nil {
+				return fatalf("%v", err)
+			}
 			recipientStrings = append(recipientStrings, s.Args[0])
 		case "add-identity":
-			expectStanzaWithNoBody(s, 1)
+			if err := expectStanzaWithNoBody(s, 1); err != nil {
+				return fatalf("%v", err)
+			}
 			identityStrings = append(identityStrings, s.Args[0])
 		case "extension-labels":
-			expectStanzaWithNoBody(s, 0)
+			if err := expectStanzaWithNoBody(s, 0); err != nil {
+				return fatalf("%v", err)
+			}
 			supportsLabels = true
 		case "wrap-file-key":
-			expectStanzaWithBody(s, 0)
+			if err := expectStanzaWithBody(s, 0); err != nil {
+				return fatalf("%v", err)
+			}
 			fileKeys = append(fileKeys, s.Body)
 		case "done":
-			expectStanzaWithNoBody(s, 0)
+			if err := expectStanzaWithNoBody(s, 0); err != nil {
+				return fatalf("%v", err)
+			}
 			break ReadLoop
 		default:
 			// Unsupported stanzas in uni-directional phases are ignored.
@@ -111,38 +124,44 @@ ReadLoop:
 	}
 
 	if len(recipientStrings)+len(identityStrings) == 0 {
-		fatalf("no recipients or identities provided")
+		return fatalf("no recipients or identities provided")
 	}
 	if len(fileKeys) == 0 {
-		fatalf("no file keys provided")
+		return fatalf("no file keys provided")
 	}
 
 	var recipients, identities []age.Recipient
 	for i, s := range recipientStrings {
 		name, data, err := ParseRecipient(s)
 		if err != nil {
-			recipientError(sr, i, err)
+			return recipientError(sr, i, err)
 		}
 		if name != p.name {
-			recipientError(sr, i, fmt.Errorf("unsupported plugin name: %q", name))
+			return recipientError(sr, i, fmt.Errorf("unsupported plugin name: %q", name))
+		}
+		if p.recipient == nil {
+			return recipientError(sr, i, fmt.Errorf("recipient encodings not supported"))
 		}
 		r, err := p.recipient(data)
 		if err != nil {
-			recipientError(sr, i, err)
+			return recipientError(sr, i, err)
 		}
 		recipients = append(recipients, r)
 	}
 	for i, s := range identityStrings {
 		name, data, err := ParseIdentity(s)
 		if err != nil {
-			identityError(sr, i, err)
+			return identityError(sr, i, err)
 		}
 		if name != p.name {
-			identityError(sr, i, fmt.Errorf("unsupported plugin name: %q", name))
+			return identityError(sr, i, fmt.Errorf("unsupported plugin name: %q", name))
+		}
+		if p.idAsRecipient == nil {
+			return identityError(sr, i, fmt.Errorf("identity encodings not supported"))
 		}
 		r, err := p.idAsRecipient(data)
 		if err != nil {
-			identityError(sr, i, err)
+			return identityError(sr, i, err)
 		}
 		identities = append(identities, r)
 	}
@@ -156,24 +175,24 @@ ReadLoop:
 		for j, r := range recipients {
 			ss, ll, err := wrapWithLabels(r, fk)
 			if err != nil {
-				recipientError(sr, j, err)
+				return recipientError(sr, j, err)
 			}
 			if i == 0 && j == 0 {
 				labels = ll
-			} else if !slicesEqual(labels, ll) {
-				recipientError(sr, j, fmt.Errorf("labels %q do not match previous recipients %q", ll, labels))
+			} else if err := checkLabels(ll, labels); err != nil {
+				return recipientError(sr, j, err)
 			}
 			stanzas[i] = append(stanzas[i], ss...)
 		}
 		for j, r := range identities {
 			ss, ll, err := wrapWithLabels(r, fk)
 			if err != nil {
-				identityError(sr, j, err)
+				return identityError(sr, j, err)
 			}
 			if i == 0 && j == 0 && len(recipients) == 0 {
 				labels = ll
-			} else if !slicesEqual(labels, ll) {
-				identityError(sr, j, fmt.Errorf("labels %q do not match previous recipients %q", ll, labels))
+			} else if err := checkLabels(ll, labels); err != nil {
+				return identityError(sr, j, err)
 			}
 			stanzas[i] = append(stanzas[i], ss...)
 		}
@@ -181,9 +200,11 @@ ReadLoop:
 
 	if supportsLabels {
 		if err := writeStanza(os.Stdout, "labels", labels...); err != nil {
-			fatalf("failed to write labels stanza: %v", err)
+			return fatalf("failed to write labels stanza: %v", err)
+		}
+		if err := expectOk(sr); err != nil {
+			return fatalf("%v", err)
 		}
-		expectOk(sr)
 	}
 
 	for i, ss := range stanzas {
@@ -191,15 +212,18 @@ ReadLoop:
 			if err := (&format.Stanza{Type: "recipient-stanza",
 				Args: append([]string{fmt.Sprint(i), s.Type}, s.Args...),
 				Body: s.Body}).Marshal(os.Stdout); err != nil {
-				fatalf("failed to write recipient-stanza: %v", err)
+				return fatalf("failed to write recipient-stanza: %v", err)
+			}
+			if err := expectOk(sr); err != nil {
+				return fatalf("%v", err)
 			}
-			expectOk(sr)
 		}
 	}
 
 	if err := writeStanza(os.Stdout, "done"); err != nil {
-		fatalf("failed to write done stanza: %v", err)
+		return fatalf("failed to write done stanza: %v", err)
 	}
+	return 0
 }
 
 func wrapWithLabels(r age.Recipient, fileKey []byte) ([]*age.Stanza, []string, error) {
@@ -210,67 +234,181 @@ func wrapWithLabels(r age.Recipient, fileKey []byte) ([]*age.Stanza, []string, e
 	return s, nil, err
 }
 
-func (p *Plugin) IdentityV1() {
-	if p.identity == nil {
-		fatalf("identity-v1 not supported")
+func checkLabels(ll, labels []string) error {
+	if !slicesEqual(ll, labels) {
+		return fmt.Errorf("labels %q do not match previous recipients %q", ll, labels)
 	}
-	panic("not implemented")
+	return nil
 }
 
-func expectStanzaWithNoBody(s *format.Stanza, wantArgs int) {
+func (p *Plugin) IdentityV1() int {
+	if p.identity == nil {
+		return fatalf("identity-v1 not supported")
+	}
+
+	var files [][]*age.Stanza
+	var identityStrings []string
+
+	sr := format.NewStanzaReader(bufio.NewReader(os.Stdin))
+ReadLoop:
+	for {
+		s, err := sr.ReadStanza()
+		if err != nil {
+			return fatalf("failed to read stanza: %v", err)
+		}
+
+		switch s.Type {
+		case "add-identity":
+			if err := expectStanzaWithNoBody(s, 1); err != nil {
+				return fatalf("%v", err)
+			}
+			identityStrings = append(identityStrings, s.Args[0])
+		case "recipient-stanza":
+			if len(s.Args) < 2 {
+				return fatalf("recipient-stanza stanza has %d arguments, want >=2", len(s.Args))
+			}
+			i, err := strconv.Atoi(s.Args[0])
+			if err != nil {
+				return fatalf("failed to parse recipient-stanza stanza argument: %v", err)
+			}
+			ss := &age.Stanza{Type: s.Args[1], Args: s.Args[2:], Body: s.Body}
+			switch i {
+			case len(files):
+				files = append(files, []*age.Stanza{ss})
+			case len(files) - 1:
+				files[len(files)-1] = append(files[len(files)-1], ss)
+			default:
+				return fatalf("unexpected file index %d, previous was %d", i, len(files)-1)
+			}
+		case "done":
+			if err := expectStanzaWithNoBody(s, 0); err != nil {
+				return fatalf("%v", err)
+			}
+			break ReadLoop
+		default:
+			// Unsupported stanzas in uni-directional phases are ignored.
+		}
+	}
+
+	if len(identityStrings) == 0 {
+		return fatalf("no identities provided")
+	}
+	if len(files) == 0 {
+		return fatalf("no stanzas provided")
+	}
+
+	var identities []age.Identity
+	for i, s := range identityStrings {
+		name, data, err := ParseIdentity(s)
+		if err != nil {
+			return identityError(sr, i, err)
+		}
+		if name != p.name {
+			return identityError(sr, i, fmt.Errorf("unsupported plugin name: %q", name))
+		}
+		if p.identity == nil {
+			return identityError(sr, i, fmt.Errorf("identity encodings not supported"))
+		}
+		r, err := p.identity(data)
+		if err != nil {
+			return identityError(sr, i, err)
+		}
+		identities = append(identities, r)
+	}
+
+	for i, ss := range files {
+		// TODO: there should be a mechanism to let the plugin decide the order
+		// in which identities are tried.
+		for _, id := range identities {
+			fk, err := id.Unwrap(ss)
+			if errors.Is(err, age.ErrIncorrectIdentity) {
+				continue
+			} else if err != nil {
+				if err := writeError(sr, []string{"stanza", fmt.Sprint(i), "0"}, err); err != nil {
+					return fatalf("%v", err)
+				}
+				// Note that we don't exit here, as the protocol allows
+				// continuing with other files.
+				break
+			}
+
+			s := &format.Stanza{Type: "file-key", Args: []string{fmt.Sprint(i)}, Body: fk}
+			if err := s.Marshal(os.Stdout); err != nil {
+				return fatalf("failed to write file-key: %v", err)
+			}
+			if err := expectOk(sr); err != nil {
+				return fatalf("%v", err)
+			}
+			break
+		}
+	}
+
+	if err := writeStanza(os.Stdout, "done"); err != nil {
+		return fatalf("failed to write done stanza: %v", err)
+	}
+	return 0
+}
+
+func expectStanzaWithNoBody(s *format.Stanza, wantArgs int) error {
 	if len(s.Args) != wantArgs {
-		fatalf("%s stanza has %d arguments, want %d", s.Type, len(s.Args), wantArgs)
+		return fmt.Errorf("%s stanza has %d arguments, want %d", s.Type, len(s.Args), wantArgs)
 	}
 	if len(s.Body) != 0 {
-		fatalf("%s stanza has %d bytes of body, want 0", s.Type, len(s.Body))
+		return fmt.Errorf("%s stanza has %d bytes of body, want 0", s.Type, len(s.Body))
 	}
+	return nil
 }
 
-func expectStanzaWithBody(s *format.Stanza, wantArgs int) {
+func expectStanzaWithBody(s *format.Stanza, wantArgs int) error {
 	if len(s.Args) != wantArgs {
-		fatalf("%s stanza has %d arguments, want %d", s.Type, len(s.Args), wantArgs)
+		return fmt.Errorf("%s stanza has %d arguments, want %d", s.Type, len(s.Args), wantArgs)
 	}
 	if len(s.Body) == 0 {
-		fatalf("%s stanza has 0 bytes of body, want >0", s.Type)
+		return fmt.Errorf("%s stanza has 0 bytes of body, want >0", s.Type)
 	}
+	return nil
 }
 
-func recipientError(sr *format.StanzaReader, idx int, err error) {
-	protocolError(sr, []string{"recipient", fmt.Sprint(idx)}, err)
+func recipientError(sr *format.StanzaReader, idx int, err error) int {
+	if err := writeError(sr, []string{"recipient", fmt.Sprint(idx)}, err); err != nil {
+		return fatalf("%v", err)
+	}
+	return 3
 }
 
-func identityError(sr *format.StanzaReader, idx int, err error) {
-	protocolError(sr, []string{"identity", fmt.Sprint(idx)}, err)
+func identityError(sr *format.StanzaReader, idx int, err error) int {
+	if err := writeError(sr, []string{"identity", fmt.Sprint(idx)}, err); err != nil {
+		return fatalf("%v", err)
+	}
+	return 3
 }
 
-func internalError(sr *format.StanzaReader, err error) {
-	protocolError(sr, []string{"internal"}, err)
+func expectOk(sr *format.StanzaReader) error {
+	ok, err := sr.ReadStanza()
+	if err != nil {
+		return fmt.Errorf("failed to read OK stanza: %v", err)
+	}
+	if ok.Type != "ok" {
+		return fmt.Errorf("expected OK stanza, got %q", ok.Type)
+	}
+	return expectStanzaWithNoBody(ok, 0)
 }
 
-func protocolError(sr *format.StanzaReader, args []string, err error) {
+func writeError(sr *format.StanzaReader, args []string, err error) error {
 	s := &format.Stanza{Type: "error", Args: args}
 	s.Body = []byte(err.Error())
 	if err := s.Marshal(os.Stdout); err != nil {
-		fatalf("failed to write error stanza: %v", err)
+		return fmt.Errorf("failed to write error stanza: %v", err)
 	}
-	expectOk(sr)
-	os.Exit(3)
+	if err := expectOk(sr); err != nil {
+		return fmt.Errorf("%v", err)
+	}
+	return nil
 }
 
-func expectOk(sr *format.StanzaReader) {
-	ok, err := sr.ReadStanza()
-	if err != nil {
-		fatalf("failed to read OK stanza: %v", err)
-	}
-	if ok.Type != "ok" {
-		fatalf("expected OK stanza, got %q", ok.Type)
-	}
-	expectStanzaWithNoBody(ok, 0)
-}
-
-func fatalf(format string, args ...interface{}) {
+func fatalf(format string, args ...interface{}) int {
 	fmt.Fprintf(os.Stderr, format, args...)
-	os.Exit(1)
+	return 1
 }
 
 func slicesEqual(s1, s2 []string) bool {