msgio.go 1.58 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"

	msgio "gx/QmXMkWRYvXJ2NTzHMEz5794997kNPGsVX8EhFiXXpoKDZr/go-msgio"
)

var Args ArgType

type ArgType struct {
	Command string
	Args    []string
}

func (a *ArgType) Arg(i int) string {
	n := i + 1
	if len(a.Args) < n {
		die(fmt.Sprintf("expected %d argument(s)", n))
	}
	return a.Args[i]
}

var usageStr = `
msgio - tool to wrap messages with msgio header

Usage
    msgio header 1020 >header
    cat file | msgio wrap >wrapped

Commands
    header <size>   output a msgio header of given size
    wrap            wrap incoming stream with msgio
`

func usage() {
	fmt.Println(strings.TrimSpace(usageStr))
	os.Exit(0)
}

func die(err string) {
	fmt.Fprintf(os.Stderr, "error: %s\n", err)
	os.Exit(-1)
}

func main() {
	if err := run(); err != nil {
		die(err.Error())
	}
}

func argParse() {
	flag.Usage = usage
	flag.Parse()

	args := flag.Args()
	if l := len(args); l < 1 || l > 2 {
		usage()
	}

	Args.Command = flag.Args()[0]
	Args.Args = flag.Args()[1:]
}

func run() error {
	argParse()

	w := os.Stdout
	r := os.Stdin

	switch Args.Command {
	case "header":
		size, err := strconv.Atoi(Args.Arg(0))
		if err != nil {
			return err
		}
		return header(w, size)
	case "wrap":
		return wrap(w, r)
	default:
		usage()
		return nil
	}
}

func header(w io.Writer, size int) error {
	return msgio.WriteLen(w, size)
}

func wrap(w io.Writer, r io.Reader) error {
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	if err := msgio.WriteLen(w, len(buf)); err != nil {
		return err
	}

	_, err = w.Write(buf)
	return err
}