From 449d72210666a2a0ef0419fb7da144214d1c0218 Mon Sep 17 00:00:00 2001 From: Filippo Valsorda Date: Tue, 18 Jun 2024 11:37:17 +0200 Subject: [PATCH] plugin: implement identities in framework Updates #485 --- cmd/age/age_test.go | 64 ++++------ plugin/client_test.go | 5 +- plugin/plugin.go | 266 ++++++++++++++++++++++++++++++++---------- 3 files changed, 226 insertions(+), 109 deletions(-) diff --git a/cmd/age/age_test.go b/cmd/age/age_test.go index 9291829..3541232 100644 --- a/cmd/age/age_test.go +++ b/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", diff --git a/plugin/client_test.go b/plugin/client_test.go index 03ddab0..ef08782 100644 --- a/plugin/client_test.go +++ b/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()) } diff --git a/plugin/plugin.go b/plugin/plugin.go index 6dde450..e246f3f 100644 --- a/plugin/plugin.go +++ b/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 {