From 6ebddc7ee311f05faedad5427159c61029ba9920 Mon Sep 17 00:00:00 2001 From: Filippo Valsorda Date: Mon, 17 Jun 2024 22:45:39 +0200 Subject: [PATCH] plugin: add plugin framework (recipient) Updates #485 --- plugin/plugin.go | 246 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 plugin/plugin.go diff --git a/plugin/plugin.go b/plugin/plugin.go new file mode 100644 index 0000000..4353b26 --- /dev/null +++ b/plugin/plugin.go @@ -0,0 +1,246 @@ +package plugin + +import ( + "bufio" + "flag" + "fmt" + "os" + + "filippo.io/age" + "filippo.io/age/internal/format" +) + +type Plugin struct { + name string + fs *flag.FlagSet + sm *string + + recipient func([]byte) (age.Recipient, error) + idAsRecipient func([]byte) (age.Recipient, error) + identity func([]byte) (age.Identity, error) +} + +func New(name string) (*Plugin, error) { + return &Plugin{name: name}, nil +} + +func (p *Plugin) Name() string { + return p.name +} + +func (p *Plugin) RegisterFlags(fs *flag.FlagSet) { + if fs == nil { + fs = flag.CommandLine + } + p.fs = fs + p.sm = fs.String("age-plugin", "", "age-plugin state machine") +} + +func (p *Plugin) HandleRecipient(f func(data []byte) (age.Recipient, error)) { + if p.recipient != nil { + panic("HandleRecipient called twice") + } + p.recipient = f +} + +func (p *Plugin) HandleIdentityAsRecipient(f func(data []byte) (age.Recipient, error)) { + if p.idAsRecipient != nil { + panic("HandleIdentityAsRecipient called twice") + } + p.idAsRecipient = f +} + +func (p *Plugin) HandleIdentity(f func(data []byte) (age.Identity, error)) { + if p.identity != nil { + panic("HandleIdentity called twice") + } + p.identity = f +} + +func (p *Plugin) Main() { + if p.fs == nil { + p.RegisterFlags(nil) + } + if !p.fs.Parsed() { + p.fs.Parse(os.Args[1:]) + } + if *p.sm == "recipient-v1" { + p.RecipientV1() + } + if *p.sm == "identity-v1" { + p.IdentityV1() + } +} + +func (p *Plugin) RecipientV1() { + if p.recipient == nil { + fatalf("recipient-v1 not supported") + } + + var recipientStrings, identityStrings []string + var fileKeys [][]byte + var supportsLabels bool + + sr := format.NewStanzaReader(bufio.NewReader(os.Stdin)) +ReadLoop: + for { + s, err := sr.ReadStanza() + if err != nil { + fatalf("failed to read stanza: %v", err) + } + + switch s.Type { + case "add-recipient": + expectStanzaWithNoBody(s, 1) + recipientStrings = append(recipientStrings, s.Args[0]) + case "add-identity": + expectStanzaWithNoBody(s, 1) + identityStrings = append(identityStrings, s.Args[0]) + case "extension-labels": + expectStanzaWithNoBody(s, 0) + supportsLabels = true + case "wrap-file-key": + expectStanzaWithBody(s, 0) + fileKeys = append(fileKeys, s.Body) + case "done": + expectStanzaWithNoBody(s, 0) + break ReadLoop + default: + // Unsupported stanzas in uni-directional phases are ignored. + } + } + + if len(recipientStrings)+len(identityStrings) == 0 { + fatalf("no recipients or identities provided") + } + if len(fileKeys) == 0 { + 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) + } + if name != p.name { + recipientError(sr, i, fmt.Errorf("unsupported plugin name: %q", name)) + } + r, err := p.recipient(data) + if err != nil { + 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) + } + if name != p.name { + identityError(sr, i, fmt.Errorf("unsupported plugin name: %q", name)) + } + r, err := p.idAsRecipient(data) + if err != nil { + identityError(sr, i, err) + } + identities = append(identities, r) + } + + stanzas := make([][]*age.Stanza, len(fileKeys)) + for i, fk := range fileKeys { + for j, r := range recipients { + ss, err := r.Wrap(fk) + if err != nil { + recipientError(sr, j, err) + } + stanzas[i] = append(stanzas[i], ss...) + } + for j, r := range identities { + ss, err := r.Wrap(fk) + if err != nil { + identityError(sr, j, err) + } + stanzas[i] = append(stanzas[i], ss...) + } + } + _ = supportsLabels // TODO + + for i, ss := range stanzas { + for _, s := range ss { + 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) + } + expectOk(sr) + } + } + + if err := writeStanza(os.Stdout, "done"); err != nil { + fatalf("failed to write done stanza: %v", err) + } +} + +func (p *Plugin) IdentityV1() { + if p.identity == nil { + fatalf("identity-v1 not supported") + } + panic("not implemented") +} + +func expectStanzaWithNoBody(s *format.Stanza, wantArgs int) { + if len(s.Args) != wantArgs { + fatalf("%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)) + } +} + +func expectStanzaWithBody(s *format.Stanza, wantArgs int) { + if len(s.Args) != wantArgs { + fatalf("%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) + } +} + +func recipientError(sr *format.StanzaReader, idx int, err error) { + protocolError(sr, []string{"recipient", fmt.Sprint(idx)}, err) +} + +func identityError(sr *format.StanzaReader, idx int, err error) { + protocolError(sr, []string{"identity", fmt.Sprint(idx)}, err) +} + +func internalError(sr *format.StanzaReader, err error) { + protocolError(sr, []string{"internal"}, err) +} + +func protocolError(sr *format.StanzaReader, args []string, err 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) + } + expectOk(sr) + os.Exit(3) +} + +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{}) { + fmt.Fprintf(os.Stderr, format, args...) + os.Exit(1) +}