1package main23import (4 "bufio"5 "context"6 "encoding/base64"7 "encoding/json"8 "errors"9 "flag"10 "fmt"11 "io"12 "os"13 "sync"14 "time"1516 "github.com/FelineStateMachine/tinyrelay/protocol/nostr"17 "github.com/FelineStateMachine/tinyrelay/tinyagent/client"18)1920type request struct {21 ID json.RawMessage `json:"id"`22 Method string `json:"method"`23 Params json.RawMessage `json:"params"`24}25type response struct {26 ID json.RawMessage `json:"id,omitempty"`27 Result any `json:"result,omitempty"`28 Error *rpcError `json:"error,omitempty"`29 Event string `json:"event,omitempty"`30 Subscription string `json:"subscription,omitempty"`31 Data any `json:"data,omitempty"`32}33type rpcError struct {34 Message string `json:"message"`35}36type rpcServer struct {37 client *client.Client38 outMu sync.Mutex39 subMu sync.Mutex40 subs map[string]context.CancelFunc41 sem chan struct{}42 wg sync.WaitGroup43 subWG sync.WaitGroup44}4546func main() {47 if err := run(os.Args[1:], os.Stdin, os.Stdout, os.Stderr); err != nil {48 fmt.Fprintln(os.Stderr, err)49 os.Exit(1)50 }51}52func run(args []string, in io.Reader, out, stderr io.Writer) error {53 if len(args) == 0 {54 return errors.New("usage: tinyagent keygen | rpc | call")55 }56 switch args[0] {57 case "keygen":58 return keygen(out)59 case "rpc":60 return rpc(args[1:], in, out)61 case "call":62 return call(args[1:], out)63 default:64 return fmt.Errorf("unknown command %q", args[0])65 }66}67func keygen(out io.Writer) error {68 secret, err := nostr.GenerateKey()69 if err != nil {70 return err71 }72 pub, err := nostr.PublicKey(secret)73 if err != nil {74 return err75 }76 return json.NewEncoder(out).Encode(map[string]string{"secret": secret, "pubkey": pub})77}78func rpc(args []string, in io.Reader, out io.Writer) error {79 fs := flag.NewFlagSet("rpc", flag.ContinueOnError)80 fs.SetOutput(io.Discard)81 relay := fs.String("relay", "", "relay URL")82 keyEnv := fs.String("key-env", "TINY_PRIVATE_KEY", "private key environment variable")83 if err := fs.Parse(args); err != nil {84 return err85 }86 if *relay == "" {87 return errors.New("--relay is required")88 }89 secret := os.Getenv(*keyEnv)90 if secret == "" {91 return fmt.Errorf("%s is not set", *keyEnv)92 }93 c, err := client.New(*relay, secret)94 if err != nil {95 return err96 }97 s := &rpcServer{client: c, subs: make(map[string]context.CancelFunc), sem: make(chan struct{}, 32)}98 return s.serve(in, out)99}100func (s *rpcServer) serve(in io.Reader, out io.Writer) error {101 scan := bufio.NewScanner(in)102 // A 32 MiB binary upload expands to roughly 43 MiB in base64, plus JSON103 // framing and the request envelope.104 scan.Buffer(make([]byte, 4096), 48<<20)105 ctx, cancel := context.WithCancel(context.Background())106 defer cancel()107 for scan.Scan() {108 var req request109 if err := json.Unmarshal(scan.Bytes(), &req); err != nil {110 s.write(out, response{Error: &rpcError{Message: "invalid request: " + err.Error()}})111 continue112 }113 if req.Method != "subscribe" {114 s.sem <- struct{}{}115 }116 s.wg.Add(1)117 go func() {118 defer s.wg.Done()119 if req.Method != "subscribe" {120 defer func() { <-s.sem }()121 }122 s.handle(ctx, req, out)123 }()124 }125 s.wg.Wait()126 s.stopAll()127 s.subWG.Wait()128 return scan.Err()129}130func (s *rpcServer) handle(parent context.Context, req request, out io.Writer) {131 if req.Method == "subscribe" {132 if err := s.launchSubscribe(parent, req.Params, out); err != nil {133 s.write(out, response{ID: req.ID, Error: &rpcError{Message: err.Error()}})134 return135 }136 s.write(out, response{ID: req.ID, Result: map[string]bool{"subscribed": true}})137 return138 }139 ctx := parent140 cancel := func() {}141 if req.Method != "subscribe" {142 ctx, cancel = context.WithTimeout(parent, 2*time.Minute)143 }144 defer cancel()145 result, err := s.dispatch(ctx, req.Method, req.Params, out)146 if err != nil {147 s.write(out, response{ID: req.ID, Error: &rpcError{Message: err.Error()}})148 return149 }150 s.write(out, response{ID: req.ID, Result: result})151}152func (s *rpcServer) dispatch(ctx context.Context, method string, raw json.RawMessage, out io.Writer) (any, error) {153 var p map[string]json.RawMessage154 if len(raw) > 0 {155 if err := json.Unmarshal(raw, &p); err != nil {156 return nil, err157 }158 }159 switch method {160 case "identity":161 return s.client.Identity(), nil162 case "publish":163 var e nostr.Event164 if err := json.Unmarshal(p["event"], &e); err != nil {165 return nil, err166 }167 return s.client.Publish(ctx, e)168 case "query":169 var f []client.Filter170 if err := json.Unmarshal(p["filter"], &f); err != nil {171 var one client.Filter172 if objectErr := json.Unmarshal(p["filter"], &one); objectErr != nil {173 return nil, err174 }175 f = []client.Filter{one}176 }177 return s.client.Query(ctx, f)178 case "request":179 var path, method string180 if err := json.Unmarshal(p["path"], &path); err != nil || path == "" {181 return nil, errors.New("path is required")182 }183 if len(p["method"]) > 0 {184 if err := json.Unmarshal(p["method"], &method); err != nil {185 return nil, err186 }187 }188 return s.client.Request(ctx, path, method, p["body"])189 case "upload":190 var room, filename, typ string191 if err := json.Unmarshal(p["room"], &room); err != nil || room == "" {192 return nil, errors.New("room is required")193 }194 if len(p["filename"]) > 0 {195 if err := json.Unmarshal(p["filename"], &filename); err != nil {196 return nil, err197 }198 }199 if len(p["type"]) > 0 {200 if err := json.Unmarshal(p["type"], &typ); err != nil {201 return nil, err202 }203 }204 var encoded string205 if err := json.Unmarshal(p["data"], &encoded); err != nil {206 return nil, err207 }208 data, err := decode(encoded)209 if err != nil {210 return nil, err211 }212 return s.client.Upload(ctx, room, filename, typ, data)213 case "download":214 var path string215 if err := json.Unmarshal(p["path"], &path); err != nil {216 return nil, err217 }218 return s.client.Download(ctx, path)219 case "unsubscribe":220 var name string221 if err := json.Unmarshal(p["subscription"], &name); err != nil {222 return nil, err223 }224 s.subMu.Lock()225 stop := s.subs[name]226 delete(s.subs, name)227 s.subMu.Unlock()228 if stop != nil {229 stop()230 }231 return map[string]bool{"stopped": stop != nil}, nil232 default:233 return nil, fmt.Errorf("unknown method %q", method)234 }235}236func (s *rpcServer) launchSubscribe(parent context.Context, raw json.RawMessage, out io.Writer) error {237 var p map[string]json.RawMessage238 if err := json.Unmarshal(raw, &p); err != nil {239 return err240 }241 var name string242 if err := json.Unmarshal(p["subscription"], &name); err != nil || name == "" {243 return errors.New("subscription is required")244 }245 var f client.Filter246 if err := json.Unmarshal(p["filter"], &f); err != nil {247 return err248 }249 ctx, cancel := context.WithCancel(parent)250 s.subMu.Lock()251 previous := s.subs[name]252 s.subs[name] = cancel253 s.subMu.Unlock()254 if previous != nil {255 previous()256 }257 s.subWG.Add(1)258 go func() {259 defer s.subWG.Done()260 err := s.client.Subscribe(ctx, name, f, func(kind string, e nostr.Event) {261 if kind == "connected" {262 s.write(out, response{Event: kind, Subscription: name})263 return264 }265 s.write(out, response{Event: "event", Subscription: kind, Data: e})266 })267 if err != nil && parent.Err() == nil {268 s.write(out, response{Event: "error", Subscription: name, Data: map[string]string{"message": err.Error()}})269 }270 // Keep the cancellation handle until an explicit unsubscribe. A newer271 // subscription may already own this name when the old stream exits.272 }()273 return nil274}275func (s *rpcServer) subscribe(parent context.Context, p map[string]json.RawMessage, out io.Writer) error {276 var name string277 if err := json.Unmarshal(p["subscription"], &name); err != nil || name == "" {278 return errors.New("subscription is required")279 }280 var f client.Filter281 if err := json.Unmarshal(p["filter"], &f); err != nil {282 return err283 }284 ctx, cancel := context.WithCancel(parent)285 s.subMu.Lock()286 s.subs[name] = cancel287 s.subMu.Unlock()288 return s.client.Subscribe(ctx, name, f, func(kind string, e nostr.Event) {289 if kind == "connected" {290 s.write(out, response{Event: kind, Subscription: name})291 return292 }293 s.write(out, response{Event: "event", Subscription: kind, Data: e})294 })295}296func (s *rpcServer) stopAll() {297 s.subMu.Lock()298 defer s.subMu.Unlock()299 for name, stop := range s.subs {300 stop()301 delete(s.subs, name)302 }303}304func (s *rpcServer) write(out io.Writer, v response) {305 s.outMu.Lock()306 defer s.outMu.Unlock()307 _ = json.NewEncoder(out).Encode(v)308}309func decode(raw string) ([]byte, error) {310 const max = 32 << 20311 if len(raw) > max*2 {312 return nil, errors.New("invalid attachment data")313 }314 b, err := base64.StdEncoding.DecodeString(raw)315 if err != nil || len(b) > max {316 return nil, errors.New("invalid attachment data")317 }318 return b, nil319}320func call(args []string, out io.Writer) error {321 fs := flag.NewFlagSet("call", flag.ContinueOnError)322 fs.SetOutput(io.Discard)323 relay := fs.String("relay", "", "relay URL")324 keyEnv := fs.String("key-env", "TINY_PRIVATE_KEY", "private key environment variable")325 if err := fs.Parse(args); err != nil {326 return err327 }328 if *relay == "" || fs.NArg() < 1 {329 return errors.New("usage: tinyagent call --relay URL METHOD [JSON_PARAMS]")330 }331 secret := os.Getenv(*keyEnv)332 if secret == "" {333 return fmt.Errorf("%s is not set", *keyEnv)334 }335 c, err := client.New(*relay, secret)336 if err != nil {337 return err338 }339 method := fs.Arg(0)340 raw := json.RawMessage(`{}`)341 if fs.NArg() > 1 {342 raw = json.RawMessage(fs.Arg(1))343 }344 server := &rpcServer{client: c, subs: make(map[string]context.CancelFunc), sem: make(chan struct{}, 1)}345 result, err := server.dispatch(context.Background(), method, raw, out)346 if err != nil {347 return err348 }349 return json.NewEncoder(out).Encode(result)350}