prehistoric-systems/age
1
0
Fork 0
mirror of https://github.com/FiloSottile/age.git synced 2026-03-11 08:55:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

plugin: add Plugin.MainWithIO

Browse source
This commit is contained in:
Filippo Valsorda 2024-09-26 12:40:18 +02:00
parent 582848919a
commit 293f859a2c
1 changed files with 93 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

172
plugin/plugin.go
View file

@ -5,6 +5,7 @@ import (
"errors" "errors"
"flag" "flag"
"fmt" "fmt"
"io"
"os" "os"
"strconv" "strconv"
@ -26,6 +27,9 @@ type Plugin struct {
idAsRecipient func([]byte) (age.Recipient, error) idAsRecipient func([]byte) (age.Recipient, error)
identity func([]byte) (age.Identity, error) identity func([]byte) (age.Identity, error)
stdin io.Reader
stdout, stderr io.Writer
sr *format.StanzaReader sr *format.StanzaReader
// broken is set if the protocol broke down during an interaction function // broken is set if the protocol broke down during an interaction function
// called by a Recipient or Identity. // called by a Recipient or Identity.
@ -107,6 +111,15 @@ func (p *Plugin) HandleIdentity(f func(data []byte) (age.Identity, error)) {
// It automatically calls [Plugin.RegisterFlags] and [flag.Parse] if they were // It automatically calls [Plugin.RegisterFlags] and [flag.Parse] if they were
// not called before. // not called before.
func (p *Plugin) Main() int { func (p *Plugin) Main() int {
return p.MainWithIO(os.Stdin, os.Stdout, os.Stderr)
}
// MainWithIO works like [Plugin.Main] but runs the plugin protocol over the
// given io.Reader and io.Writers.
func (p *Plugin) MainWithIO(stdin io.Reader, stdout, stderr io.Writer) int {
p.stdin = stdin
p.stdout = stdout
p.stderr = stderr
if p.fs == nil { if p.fs == nil {
p.RegisterFlags(nil) p.RegisterFlags(nil)
} }
@ -119,7 +132,8 @@ func (p *Plugin) Main() int {
if *p.sm == "identity-v1" { if *p.sm == "identity-v1" {
return p.IdentityV1() return p.IdentityV1()
} }
return fatalf("unknown state machine %q", *p.sm) fmt.Fprintf(p.stderr, "unknown state machine %q", *p.sm)
return 4
} }
// RecipientV1 implements the recipient-v1 state machine over stdin/stdout, and // RecipientV1 implements the recipient-v1 state machine over stdin/stdout, and
@ -128,45 +142,45 @@ func (p *Plugin) Main() int {
// Most plugins should call [Plugin.Main] instead of this method. // Most plugins should call [Plugin.Main] instead of this method.
func (p *Plugin) RecipientV1() int { func (p *Plugin) RecipientV1() int {
if p.recipient == nil && p.idAsRecipient == nil { if p.recipient == nil && p.idAsRecipient == nil {
return fatalf("recipient-v1 not supported") return p.fatalf("recipient-v1 not supported")
} }
var recipientStrings, identityStrings []string var recipientStrings, identityStrings []string
var fileKeys [][]byte var fileKeys [][]byte
var supportsLabels bool var supportsLabels bool
p.sr = format.NewStanzaReader(bufio.NewReader(os.Stdin)) p.sr = format.NewStanzaReader(bufio.NewReader(p.stdin))
ReadLoop: ReadLoop:
for { for {
s, err := p.sr.ReadStanza() s, err := p.sr.ReadStanza()
if err != nil { if err != nil {
return fatalf("failed to read stanza: %v", err) return p.fatalf("failed to read stanza: %v", err)
} }
switch s.Type { switch s.Type {
case "add-recipient": case "add-recipient":
if err := expectStanzaWithNoBody(s, 1); err != nil { if err := expectStanzaWithNoBody(s, 1); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
recipientStrings = append(recipientStrings, s.Args[0]) recipientStrings = append(recipientStrings, s.Args[0])
case "add-identity": case "add-identity":
if err := expectStanzaWithNoBody(s, 1); err != nil { if err := expectStanzaWithNoBody(s, 1); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
identityStrings = append(identityStrings, s.Args[0]) identityStrings = append(identityStrings, s.Args[0])
case "extension-labels": case "extension-labels":
if err := expectStanzaWithNoBody(s, 0); err != nil { if err := expectStanzaWithNoBody(s, 0); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
supportsLabels = true supportsLabels = true
case "wrap-file-key": case "wrap-file-key":
if err := expectStanzaWithBody(s, 0); err != nil { if err := expectStanzaWithBody(s, 0); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
fileKeys = append(fileKeys, s.Body) fileKeys = append(fileKeys, s.Body)
case "done": case "done":
if err := expectStanzaWithNoBody(s, 0); err != nil { if err := expectStanzaWithNoBody(s, 0); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
break ReadLoop break ReadLoop
default: default:
@ -175,44 +189,44 @@ ReadLoop:
} }
if len(recipientStrings)+len(identityStrings) == 0 { if len(recipientStrings)+len(identityStrings) == 0 {
return fatalf("no recipients or identities provided") return p.fatalf("no recipients or identities provided")
} }
if len(fileKeys) == 0 { if len(fileKeys) == 0 {
return fatalf("no file keys provided") return p.fatalf("no file keys provided")
} }
var recipients, identities []age.Recipient var recipients, identities []age.Recipient
for i, s := range recipientStrings { for i, s := range recipientStrings {
name, data, err := ParseRecipient(s) name, data, err := ParseRecipient(s)
if err != nil { if err != nil {
return recipientError(p.sr, i, err) return p.recipientError(i, err)
} }
if name != p.name { if name != p.name {
return recipientError(p.sr, i, fmt.Errorf("unsupported plugin name: %q", name)) return p.recipientError(i, fmt.Errorf("unsupported plugin name: %q", name))
} }
if p.recipient == nil { if p.recipient == nil {
return recipientError(p.sr, i, fmt.Errorf("recipient encodings not supported")) return p.recipientError(i, fmt.Errorf("recipient encodings not supported"))
} }
r, err := p.recipient(data) r, err := p.recipient(data)
if err != nil { if err != nil {
return recipientError(p.sr, i, err) return p.recipientError(i, err)
} }
recipients = append(recipients, r) recipients = append(recipients, r)
} }
for i, s := range identityStrings { for i, s := range identityStrings {
name, data, err := ParseIdentity(s) name, data, err := ParseIdentity(s)
if err != nil { if err != nil {
return identityError(p.sr, i, err) return p.identityError(i, err)
} }
if name != p.name { if name != p.name {
return identityError(p.sr, i, fmt.Errorf("unsupported plugin name: %q", name)) return p.identityError(i, fmt.Errorf("unsupported plugin name: %q", name))
} }
if p.idAsRecipient == nil { if p.idAsRecipient == nil {
return identityError(p.sr, i, fmt.Errorf("identity encodings not supported")) return p.identityError(i, fmt.Errorf("identity encodings not supported"))
} }
r, err := p.idAsRecipient(data) r, err := p.idAsRecipient(data)
if err != nil { if err != nil {
return identityError(p.sr, i, err) return p.identityError(i, err)
} }
identities = append(identities, r) identities = append(identities, r)
} }
@ -228,12 +242,12 @@ ReadLoop:
if p.broken { if p.broken {
return 2 return 2
} else if err != nil { } else if err != nil {
return recipientError(p.sr, j, err) return p.recipientError(j, err)
} }
if i == 0 && j == 0 { if i == 0 && j == 0 {
labels = ll labels = ll
} else if err := checkLabels(ll, labels); err != nil { } else if err := checkLabels(ll, labels); err != nil {
return recipientError(p.sr, j, err) return p.recipientError(j, err)
} }
stanzas[i] = append(stanzas[i], ss...) stanzas[i] = append(stanzas[i], ss...)
} }
@ -242,31 +256,31 @@ ReadLoop:
if p.broken { if p.broken {
return 2 return 2
} else if err != nil { } else if err != nil {
return identityError(p.sr, j, err) return p.identityError(j, err)
} }
if i == 0 && j == 0 && len(recipients) == 0 { if i == 0 && j == 0 && len(recipients) == 0 {
labels = ll labels = ll
} else if err := checkLabels(ll, labels); err != nil { } else if err := checkLabels(ll, labels); err != nil {
return identityError(p.sr, j, err) return p.identityError(j, err)
} }
stanzas[i] = append(stanzas[i], ss...) stanzas[i] = append(stanzas[i], ss...)
} }
} }
if sent, err := writeGrease(os.Stdout); err != nil { if sent, err := writeGrease(p.stdout); err != nil {
return fatalf("failed to write grease: %v", err) return p.fatalf("failed to write grease: %v", err)
} else if sent { } else if sent {
if err := expectUnsupported(p.sr); err != nil { if err := expectUnsupported(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
} }
if supportsLabels { if supportsLabels {
if err := writeStanza(os.Stdout, "labels", labels...); err != nil { if err := writeStanza(p.stdout, "labels", labels...); err != nil {
return fatalf("failed to write labels stanza: %v", err) return p.fatalf("failed to write labels stanza: %v", err)
} }
if err := expectOk(p.sr); err != nil { if err := expectOk(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
} }
@ -274,24 +288,24 @@ ReadLoop:
for _, s := range ss { for _, s := range ss {
if err := (&format.Stanza{Type: "recipient-stanza", if err := (&format.Stanza{Type: "recipient-stanza",
Args: append([]string{fmt.Sprint(i), s.Type}, s.Args...), Args: append([]string{fmt.Sprint(i), s.Type}, s.Args...),
Body: s.Body}).Marshal(os.Stdout); err != nil { Body: s.Body}).Marshal(p.stdout); err != nil {
return fatalf("failed to write recipient-stanza: %v", err) return p.fatalf("failed to write recipient-stanza: %v", err)
} }
if err := expectOk(p.sr); err != nil { if err := expectOk(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
} }
if sent, err := writeGrease(os.Stdout); err != nil { if sent, err := writeGrease(p.stdout); err != nil {
return fatalf("failed to write grease: %v", err) return p.fatalf("failed to write grease: %v", err)
} else if sent { } else if sent {
if err := expectUnsupported(p.sr); err != nil { if err := expectUnsupported(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
} }
} }
if err := writeStanza(os.Stdout, "done"); err != nil { if err := writeStanza(p.stdout, "done"); err != nil {
return fatalf("failed to write done stanza: %v", err) return p.fatalf("failed to write done stanza: %v", err)
} }
return 0 return 0
} }
@ -317,33 +331,33 @@ func checkLabels(ll, labels []string) error {
// Most plugins should call [Plugin.Main] instead of this method. // Most plugins should call [Plugin.Main] instead of this method.
func (p *Plugin) IdentityV1() int { func (p *Plugin) IdentityV1() int {
if p.identity == nil { if p.identity == nil {
return fatalf("identity-v1 not supported") return p.fatalf("identity-v1 not supported")
} }
var files [][]*age.Stanza var files [][]*age.Stanza
var identityStrings []string var identityStrings []string
p.sr = format.NewStanzaReader(bufio.NewReader(os.Stdin)) p.sr = format.NewStanzaReader(bufio.NewReader(p.stdin))
ReadLoop: ReadLoop:
for { for {
s, err := p.sr.ReadStanza() s, err := p.sr.ReadStanza()
if err != nil { if err != nil {
return fatalf("failed to read stanza: %v", err) return p.fatalf("failed to read stanza: %v", err)
} }
switch s.Type { switch s.Type {
case "add-identity": case "add-identity":
if err := expectStanzaWithNoBody(s, 1); err != nil { if err := expectStanzaWithNoBody(s, 1); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
identityStrings = append(identityStrings, s.Args[0]) identityStrings = append(identityStrings, s.Args[0])
case "recipient-stanza": case "recipient-stanza":
if len(s.Args) < 2 { if len(s.Args) < 2 {
return fatalf("recipient-stanza stanza has %d arguments, want >=2", len(s.Args)) return p.fatalf("recipient-stanza stanza has %d arguments, want >=2", len(s.Args))
} }
i, err := strconv.Atoi(s.Args[0]) i, err := strconv.Atoi(s.Args[0])
if err != nil { if err != nil {
return fatalf("failed to parse recipient-stanza stanza argument: %v", err) return p.fatalf("failed to parse recipient-stanza stanza argument: %v", err)
} }
ss := &age.Stanza{Type: s.Args[1], Args: s.Args[2:], Body: s.Body} ss := &age.Stanza{Type: s.Args[1], Args: s.Args[2:], Body: s.Body}
switch i { switch i {
@ -352,11 +366,11 @@ ReadLoop:
case len(files) - 1: case len(files) - 1:
files[len(files)-1] = append(files[len(files)-1], ss) files[len(files)-1] = append(files[len(files)-1], ss)
default: default:
return fatalf("unexpected file index %d, previous was %d", i, len(files)-1) return p.fatalf("unexpected file index %d, previous was %d", i, len(files)-1)
} }
case "done": case "done":
if err := expectStanzaWithNoBody(s, 0); err != nil { if err := expectStanzaWithNoBody(s, 0); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
break ReadLoop break ReadLoop
default: default:
@ -365,37 +379,37 @@ ReadLoop:
} }
if len(identityStrings) == 0 { if len(identityStrings) == 0 {
return fatalf("no identities provided") return p.fatalf("no identities provided")
} }
if len(files) == 0 { if len(files) == 0 {
return fatalf("no stanzas provided") return p.fatalf("no stanzas provided")
} }
var identities []age.Identity var identities []age.Identity
for i, s := range identityStrings { for i, s := range identityStrings {
name, data, err := ParseIdentity(s) name, data, err := ParseIdentity(s)
if err != nil { if err != nil {
return identityError(p.sr, i, err) return p.identityError(i, err)
} }
if name != p.name { if name != p.name {
return identityError(p.sr, i, fmt.Errorf("unsupported plugin name: %q", name)) return p.identityError(i, fmt.Errorf("unsupported plugin name: %q", name))
} }
if p.identity == nil { if p.identity == nil {
return identityError(p.sr, i, fmt.Errorf("identity encodings not supported")) return p.identityError(i, fmt.Errorf("identity encodings not supported"))
} }
r, err := p.identity(data) r, err := p.identity(data)
if err != nil { if err != nil {
return identityError(p.sr, i, err) return p.identityError(i, err)
} }
identities = append(identities, r) identities = append(identities, r)
} }
for i, ss := range files { for i, ss := range files {
if sent, err := writeGrease(os.Stdout); err != nil { if sent, err := writeGrease(p.stdout); err != nil {
return fatalf("failed to write grease: %v", err) return p.fatalf("failed to write grease: %v", err)
} else if sent { } else if sent {
if err := expectUnsupported(p.sr); err != nil { if err := expectUnsupported(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
} }
@ -408,8 +422,8 @@ ReadLoop:
} else if errors.Is(err, age.ErrIncorrectIdentity) { } else if errors.Is(err, age.ErrIncorrectIdentity) {
continue continue
} else if err != nil { } else if err != nil {
if err := writeError(p.sr, []string{"stanza", fmt.Sprint(i), "0"}, err); err != nil { if err := p.writeError([]string{"stanza", fmt.Sprint(i), "0"}, err); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
// Note that we don't exit here, as the protocol allows // Note that we don't exit here, as the protocol allows
// continuing with other files. // continuing with other files.
@ -417,18 +431,18 @@ ReadLoop:
} }
s := &format.Stanza{Type: "file-key", Args: []string{fmt.Sprint(i)}, Body: fk} s := &format.Stanza{Type: "file-key", Args: []string{fmt.Sprint(i)}, Body: fk}
if err := s.Marshal(os.Stdout); err != nil { if err := s.Marshal(p.stdout); err != nil {
return fatalf("failed to write file-key: %v", err) return p.fatalf("failed to write file-key: %v", err)
} }
if err := expectOk(p.sr); err != nil { if err := expectOk(p.sr); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
break break
} }
} }
if err := writeStanza(os.Stdout, "done"); err != nil { if err := writeStanza(p.stdout, "done"); err != nil {
return fatalf("failed to write done stanza: %v", err) return p.fatalf("failed to write done stanza: %v", err)
} }
return 0 return 0
} }
@ -440,7 +454,7 @@ ReadLoop:
// //
// It must only be called by a Wrap or Unwrap method invoked by [Plugin.Main]. // It must only be called by a Wrap or Unwrap method invoked by [Plugin.Main].
func (p *Plugin) DisplayMessage(message string) error { func (p *Plugin) DisplayMessage(message string) error {
if err := writeStanzaWithBody(os.Stdout, "msg", []byte(message)); err != nil { if err := writeStanzaWithBody(p.stdout, "msg", []byte(message)); err != nil {
return p.fatalInteractf("failed to write msg stanza: %v", err) return p.fatalInteractf("failed to write msg stanza: %v", err)
} }
s, err := readOkOrFail(p.sr) s, err := readOkOrFail(p.sr)
@ -466,7 +480,7 @@ func (p *Plugin) RequestValue(prompt string, secret bool) (string, error) {
if secret { if secret {
t = "request-secret" t = "request-secret"
} }
if err := writeStanzaWithBody(os.Stdout, t, []byte(prompt)); err != nil { if err := writeStanzaWithBody(p.stdout, t, []byte(prompt)); err != nil {
return "", p.fatalInteractf("failed to write stanza: %v", err) return "", p.fatalInteractf("failed to write stanza: %v", err)
} }
s, err := readOkOrFail(p.sr) s, err := readOkOrFail(p.sr)
@ -495,7 +509,7 @@ func (p *Plugin) Confirm(prompt, yes, no string) (choseYes bool, err error) {
args = append(args, format.EncodeToString([]byte(no))) args = append(args, format.EncodeToString([]byte(no)))
} }
s := &format.Stanza{Type: "confirm", Args: args, Body: []byte(prompt)} s := &format.Stanza{Type: "confirm", Args: args, Body: []byte(prompt)}
if err := s.Marshal(os.Stdout); err != nil { if err := s.Marshal(p.stdout); err != nil {
return false, p.fatalInteractf("failed to write confirm stanza: %v", err) return false, p.fatalInteractf("failed to write confirm stanza: %v", err)
} }
s, err = readOkOrFail(p.sr) s, err = readOkOrFail(p.sr)
@ -515,10 +529,15 @@ func (p *Plugin) Confirm(prompt, yes, no string) (choseYes bool, err error) {
// Wrap/Unwrap caller can exit with an error. // Wrap/Unwrap caller can exit with an error.
func (p *Plugin) fatalInteractf(format string, args ...interface{}) error { func (p *Plugin) fatalInteractf(format string, args ...interface{}) error {
p.broken = true p.broken = true
fmt.Fprintf(os.Stderr, format, args...) fmt.Fprintf(p.stderr, format, args...)
return fmt.Errorf(format, args...) return fmt.Errorf(format, args...)
} }
func (p *Plugin) fatalf(format string, args ...interface{}) int {
fmt.Fprintf(p.stderr, format, args...)
return 1
}
func expectStanzaWithNoBody(s *format.Stanza, wantArgs int) error { func expectStanzaWithNoBody(s *format.Stanza, wantArgs int) error {
if len(s.Args) != wantArgs { if len(s.Args) != wantArgs {
return fmt.Errorf("%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)
@ -539,16 +558,16 @@ func expectStanzaWithBody(s *format.Stanza, wantArgs int) error {
return nil return nil
} }
func recipientError(sr *format.StanzaReader, idx int, err error) int { func (p *Plugin) recipientError(idx int, err error) int {
if err := writeError(sr, []string{"recipient", fmt.Sprint(idx)}, err); err != nil { if err := p.writeError([]string{"recipient", fmt.Sprint(idx)}, err); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
return 3 return 3
} }
func identityError(sr *format.StanzaReader, idx int, err error) int { func (p *Plugin) identityError(idx int, err error) int {
if err := writeError(sr, []string{"identity", fmt.Sprint(idx)}, err); err != nil { if err := p.writeError([]string{"identity", fmt.Sprint(idx)}, err); err != nil {
return fatalf("%v", err) return p.fatalf("%v", err)
} }
return 3 return 3
} }
@ -593,23 +612,18 @@ func expectUnsupported(sr *format.StanzaReader) error {
return expectStanzaWithNoBody(unsupported, 0) return expectStanzaWithNoBody(unsupported, 0)
} }
func writeError(sr *format.StanzaReader, args []string, err error) error { func (p *Plugin) writeError(args []string, err error) error {
s := &format.Stanza{Type: "error", Args: args} s := &format.Stanza{Type: "error", Args: args}
s.Body = []byte(err.Error()) s.Body = []byte(err.Error())
if err := s.Marshal(os.Stdout); err != nil { if err := s.Marshal(p.stderr); err != nil {
return fmt.Errorf("failed to write error stanza: %v", err) return fmt.Errorf("failed to write error stanza: %v", err)
} }
if err := expectOk(sr); err != nil { if err := expectOk(p.sr); err != nil {
return fmt.Errorf("%v", err) return fmt.Errorf("%v", err)
} }
return nil return nil
} }
func fatalf(format string, args ...interface{}) int {
fmt.Fprintf(os.Stderr, format, args...)
return 1
}
func slicesEqual(s1, s2 []string) bool { func slicesEqual(s1, s2 []string) bool {
if len(s1) != len(s2) { if len(s1) != len(s2) {
return false return false