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(exitStatus(err))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 | mcp | diagnose")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 case "mcp":64 return mcpCommand(args[1:], in, out, stderr)65 case "diagnose":66 return diagnoseCommand(args[1:], out)67 default:68 return fmt.Errorf("unknown command %q", args[0])69 }70}71func keygen(out io.Writer) error {72 secret, err := nostr.GenerateKey()73 if err != nil {74 return err75 }76 pub, err := nostr.PublicKey(secret)77 if err != nil {78 return err79 }80 return json.NewEncoder(out).Encode(map[string]string{"secret": secret, "pubkey": pub})81}82func rpc(args []string, in io.Reader, out io.Writer) error {83 fs := flag.NewFlagSet("rpc", flag.ContinueOnError)84 fs.SetOutput(io.Discard)85 relay := fs.String("relay", "", "relay URL")86 keyEnv := fs.String("key-env", "TINY_PRIVATE_KEY", "private key environment variable")87 if err := fs.Parse(args); err != nil {88 return err89 }90 if *relay == "" {91 return errors.New("--relay is required")92 }93 secret := os.Getenv(*keyEnv)94 if secret == "" {95 return fmt.Errorf("%s is not set", *keyEnv)96 }97 c, err := client.New(*relay, secret)98 if err != nil {99 return err100 }101 s := &rpcServer{client: c, subs: make(map[string]context.CancelFunc), sem: make(chan struct{}, 32)}102 return s.serve(in, out)103}104func (s *rpcServer) serve(in io.Reader, out io.Writer) error {105 scan := bufio.NewScanner(in)106 // A 32 MiB binary upload expands to roughly 43 MiB in base64, plus JSON107 // framing and the request envelope.108 scan.Buffer(make([]byte, 4096), 48<<20)109 ctx, cancel := context.WithCancel(context.Background())110 defer cancel()111 for scan.Scan() {112 var req request113 if err := json.Unmarshal(scan.Bytes(), &req); err != nil {114 s.write(out, response{Error: &rpcError{Message: "invalid request: " + err.Error()}})115 continue116 }117 if req.Method != "subscribe" {118 s.sem <- struct{}{}119 }120 s.wg.Add(1)121 go func() {122 defer s.wg.Done()123 if req.Method != "subscribe" {124 defer func() { <-s.sem }()125 }126 s.handle(ctx, req, out)127 }()128 }129 s.wg.Wait()130 s.stopAll()131 s.subWG.Wait()132 return scan.Err()133}134func (s *rpcServer) handle(parent context.Context, req request, out io.Writer) {135 if req.Method == "subscribe" {136 if err := s.launchSubscribe(parent, req.Params, out); err != nil {137 s.write(out, response{ID: req.ID, Error: &rpcError{Message: err.Error()}})138 return139 }140 s.write(out, response{ID: req.ID, Result: map[string]bool{"subscribed": true}})141 return142 }143 ctx := parent144 cancel := func() {}145 if req.Method != "subscribe" {146 ctx, cancel = context.WithTimeout(parent, 2*time.Minute)147 }148 defer cancel()149 result, err := s.dispatch(ctx, req.Method, req.Params, out)150 if err != nil {151 s.write(out, response{ID: req.ID, Error: &rpcError{Message: err.Error()}})152 return153 }154 s.write(out, response{ID: req.ID, Result: result})155}156func (s *rpcServer) dispatch(ctx context.Context, method string, raw json.RawMessage, out io.Writer) (any, error) {157 var p map[string]json.RawMessage158 if len(raw) > 0 {159 if err := json.Unmarshal(raw, &p); err != nil {160 return nil, err161 }162 }163 switch method {164 case "identity":165 return s.client.Identity(), nil166 case "publish":167 var e nostr.Event168 if err := json.Unmarshal(p["event"], &e); err != nil {169 return nil, err170 }171 return s.client.Publish(ctx, e)172 case "query":173 var f []client.Filter174 if err := json.Unmarshal(p["filter"], &f); err != nil {175 var one client.Filter176 if objectErr := json.Unmarshal(p["filter"], &one); objectErr != nil {177 return nil, err178 }179 f = []client.Filter{one}180 }181 return s.client.Query(ctx, f)182 case "request":183 var path, method string184 if err := json.Unmarshal(p["path"], &path); err != nil || path == "" {185 return nil, errors.New("path is required")186 }187 if len(p["method"]) > 0 {188 if err := json.Unmarshal(p["method"], &method); err != nil {189 return nil, err190 }191 }192 return s.client.Request(ctx, path, method, p["body"])193 case "upload":194 var room, filename, typ string195 if err := json.Unmarshal(p["room"], &room); err != nil || room == "" {196 return nil, errors.New("room is required")197 }198 if len(p["filename"]) > 0 {199 if err := json.Unmarshal(p["filename"], &filename); err != nil {200 return nil, err201 }202 }203 if len(p["type"]) > 0 {204 if err := json.Unmarshal(p["type"], &typ); err != nil {205 return nil, err206 }207 }208 var encoded string209 if err := json.Unmarshal(p["data"], &encoded); err != nil {210 return nil, err211 }212 data, err := decode(encoded)213 if err != nil {214 return nil, err215 }216 return s.client.Upload(ctx, room, filename, typ, data)217 case "download":218 var path string219 if err := json.Unmarshal(p["path"], &path); err != nil {220 return nil, err221 }222 return s.client.Download(ctx, path)223 case "diagnose":224 return s.diagnose(ctx, p)225 case "unsubscribe":226 var name string227 if err := json.Unmarshal(p["subscription"], &name); err != nil {228 return nil, err229 }230 s.subMu.Lock()231 stop := s.subs[name]232 delete(s.subs, name)233 s.subMu.Unlock()234 if stop != nil {235 stop()236 }237 return map[string]bool{"stopped": stop != nil}, nil238 default:239 return nil, fmt.Errorf("unknown method %q", method)240 }241}242func (s *rpcServer) launchSubscribe(parent context.Context, raw json.RawMessage, out io.Writer) error {243 var p map[string]json.RawMessage244 if err := json.Unmarshal(raw, &p); err != nil {245 return err246 }247 var name string248 if err := json.Unmarshal(p["subscription"], &name); err != nil || name == "" {249 return errors.New("subscription is required")250 }251 var f client.Filter252 if err := json.Unmarshal(p["filter"], &f); err != nil {253 return err254 }255 ctx, cancel := context.WithCancel(parent)256 s.subMu.Lock()257 previous := s.subs[name]258 s.subs[name] = cancel259 s.subMu.Unlock()260 if previous != nil {261 previous()262 }263 s.subWG.Add(1)264 go func() {265 defer s.subWG.Done()266 err := s.client.Subscribe(ctx, name, f, func(kind string, e nostr.Event) {267 if kind == "connected" {268 s.write(out, response{Event: kind, Subscription: name})269 return270 }271 s.write(out, response{Event: "event", Subscription: kind, Data: e})272 })273 if err != nil && parent.Err() == nil {274 s.write(out, response{Event: "error", Subscription: name, Data: map[string]string{"message": err.Error()}})275 }276 // Keep the cancellation handle until an explicit unsubscribe. A newer277 // subscription may already own this name when the old stream exits.278 }()279 return nil280}281func (s *rpcServer) subscribe(parent context.Context, p map[string]json.RawMessage, out io.Writer) error {282 var name string283 if err := json.Unmarshal(p["subscription"], &name); err != nil || name == "" {284 return errors.New("subscription is required")285 }286 var f client.Filter287 if err := json.Unmarshal(p["filter"], &f); err != nil {288 return err289 }290 ctx, cancel := context.WithCancel(parent)291 s.subMu.Lock()292 s.subs[name] = cancel293 s.subMu.Unlock()294 return s.client.Subscribe(ctx, name, f, func(kind string, e nostr.Event) {295 if kind == "connected" {296 s.write(out, response{Event: kind, Subscription: name})297 return298 }299 s.write(out, response{Event: "event", Subscription: kind, Data: e})300 })301}302func (s *rpcServer) stopAll() {303 s.subMu.Lock()304 defer s.subMu.Unlock()305 for name, stop := range s.subs {306 stop()307 delete(s.subs, name)308 }309}310func (s *rpcServer) write(out io.Writer, v response) {311 s.outMu.Lock()312 defer s.outMu.Unlock()313 _ = json.NewEncoder(out).Encode(v)314}315func decode(raw string) ([]byte, error) {316 const max = 32 << 20317 if len(raw) > max*2 {318 return nil, errors.New("invalid attachment data")319 }320 b, err := base64.StdEncoding.DecodeString(raw)321 if err != nil || len(b) > max {322 return nil, errors.New("invalid attachment data")323 }324 return b, nil325}326func call(args []string, out io.Writer) error {327 fs := flag.NewFlagSet("call", flag.ContinueOnError)328 fs.SetOutput(io.Discard)329 relay := fs.String("relay", "", "relay URL")330 keyEnv := fs.String("key-env", "TINY_PRIVATE_KEY", "private key environment variable")331 if err := fs.Parse(args); err != nil {332 return err333 }334 if *relay == "" || fs.NArg() < 1 {335 return errors.New("usage: tinyagent call --relay URL METHOD [JSON_PARAMS]")336 }337 secret := os.Getenv(*keyEnv)338 if secret == "" {339 return fmt.Errorf("%s is not set", *keyEnv)340 }341 c, err := client.New(*relay, secret)342 if err != nil {343 return err344 }345 method := fs.Arg(0)346 raw := json.RawMessage(`{}`)347 if fs.NArg() > 1 {348 raw = json.RawMessage(fs.Arg(1))349 }350 server := &rpcServer{client: c, subs: make(map[string]context.CancelFunc), sem: make(chan struct{}, 1)}351 result, err := server.dispatch(context.Background(), method, raw, out)352 if err != nil {353 return err354 }355 return json.NewEncoder(out).Encode(result)356}