mse/cmd/mse/main.go

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