ucat.go 1.52 KB
Newer Older
Jeromy's avatar
Jeromy committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"

	"QmWsa476RGjb9scWzcRVts3QZsYjU5Kt6Y9qe8Q3vc5FHR/envpprof"

	"QmVBEPpwYNGb5xQW6jPzYgV86KSPRrzhe4hA7o7mtAHZsT/go-utp"
)

func main() {
	defer envpprof.Stop()
	listen := flag.Bool("l", false, "listen")
	port := flag.Int("p", 0, "port to listen on")
	flag.Parse()
	var (
		conn net.Conn
		err  error
	)
	if *listen {
		s, err := utp.NewSocket("udp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer s.Close()
		conn, err = s.Accept()
		if err != nil {
			log.Fatal(err)
		}
	} else {
		conn, err = utp.Dial(net.JoinHostPort(flag.Arg(0), flag.Arg(1)))
		if err != nil {
			log.Fatal(err)
		}
	}
	defer conn.Close()
	go func() {
		sig := make(chan os.Signal, 1)
		signal.Notify(sig, os.Interrupt)
		<-sig
		conn.Close()
	}()
	writerDone := make(chan struct{})
	readerDone := make(chan struct{})
	go func() {
		defer close(writerDone)
		written, err := io.Copy(conn, os.Stdin)
		if err != nil {
			conn.Close()
			log.Fatalf("error after writing %d bytes: %s", written, err)
		}
		log.Printf("wrote %d bytes", written)
		conn.Close()
	}()
	go func() {
		defer close(readerDone)
		n, err := io.Copy(os.Stdout, conn)
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("received %d bytes", n)
	}()
	// Technically we should wait until both reading and writing are done. But
	// ucat-style binaries terminate abrubtly when read or write is completed,
	// and no state remains to clean-up the peer neatly.
	select {
	case <-writerDone:
	case <-readerDone:
	}
}