mse/cmd/mse/main.go

   1 package main
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "log"
   7         "net"
   8         "os"
   9         "sync"
  10
  11         "github.com/alexflint/go-arg"
  12         "github.com/anacrolix/torrent/mse"
  13 )
  14
  15 func main() {
  16         err := mainErr()
  17         if err != nil {
  18                 log.Fatalf("fatal error: %v", err)
  19         }
  20 }
  21
  22 func mainErr() error {
  23         args := struct {
  24                 CryptoMethod mse.CryptoMethod
  25                 Dial         *struct {
  26                         Network        string `arg:"positional"`
  27                         Address        string `arg:"positional"`
  28                         SecretKey      string `arg:"positional"`
  29                         InitialPayload []byte
  30                 } `arg:"subcommand"`
  31                 Listen *struct {
  32                         Network    string   `arg:"positional"`
  33                         Address    string   `arg:"positional"`
  34                         SecretKeys []string `arg:"positional"`
  35                 } `arg:"subcommand"`
  36         }{
  37                 CryptoMethod: mse.AllSupportedCrypto,
  38         }
  39         p := arg.MustParse(&args)
  40         if args.Dial != nil {
  41                 cn, err := net.Dial(args.Dial.Network, args.Dial.Address)
  42                 if err != nil {
  43                         return fmt.Errorf("dialing: %w", err)
  44                 }
  45                 defer cn.Close()
  46                 rw, _, err := mse.InitiateHandshake(cn, []byte(args.Dial.SecretKey), args.Dial.InitialPayload, args.CryptoMethod)
  47                 if err != nil {
  48                         return fmt.Errorf("initiating handshake: %w", err)
  49                 }
  50                 doStreaming(rw)
  51         }
  52         if args.Listen != nil {
  53                 l, err := net.Listen(args.Listen.Network, args.Listen.Address)
  54                 if err != nil {
  55                         return fmt.Errorf("listening: %w", err)
  56                 }
  57                 defer l.Close()
  58                 cn, err := l.Accept()
  59                 l.Close()
  60                 if err != nil {
  61                         return fmt.Errorf("accepting: %w", err)
  62                 }
  63                 defer cn.Close()
  64                 rw, _, err := mse.ReceiveHandshake(cn, func(f func([]byte) bool) {
  65                         for _, sk := range args.Listen.SecretKeys {
  66                                 f([]byte(sk))
  67                         }
  68                 }, mse.DefaultCryptoSelector)
  69                 if err != nil {
  70                         log.Fatalf("error receiving: %v", err)
  71                 }
  72                 doStreaming(rw)
  73         }
  74         if p.Subcommand() == nil {
  75                 p.Fail("missing subcommand")
  76         }
  77         return nil
  78 }
  79
  80 func doStreaming(rw io.ReadWriter) {
  81         var wg sync.WaitGroup
  82         wg.Add(2)
  83         go func() {
  84                 defer wg.Done()
  85                 log.Println(io.Copy(rw, os.Stdin))
  86         }()
  87         go func() {
  88                 defer wg.Done()
  89                 log.Println(io.Copy(os.Stdout, rw))
  90         }()
  91         wg.Wait()
  92 }