wildcat/bufio.go at master · evanphx/wildcat · GitHub

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
package wildcat

import "io"

type sizedBodyReader struct {
	size int64
	rest []byte
	c    io.ReadCloser
}

func (br *sizedBodyReader) Read(buf []byte) (int, error) {
	if br.size == 0 {
		return 0, io.EOF
	}

	if br.rest != nil {
		if len(buf) < len(br.rest) {
			copy(buf, br.rest[:len(buf)])

			br.rest = br.rest[len(buf):]
			br.size -= int64(len(buf))
			return len(buf), nil
		} else {
			l := len(br.rest)
			copy(buf, br.rest)

			br.rest = nil
			br.size -= int64(l)
			return l, nil
		}
	}

	n, err := br.c.Read(buf[:br.size])
	if err != nil {
		return 0, err
	}

	br.size -= int64(n)
	return n, nil
}

func (br *sizedBodyReader) Close() error {
	return br.c.Close()
}

type unsizedBodyReader struct {
	rest []byte
	c    io.ReadCloser
}

func (br *unsizedBodyReader) Read(buf []byte) (int, error) {
	if br.rest != nil {
		if len(buf) < len(br.rest) {
			copy(buf, br.rest[:len(buf)])

			br.rest = br.rest[len(buf):]
			return len(buf), nil
		} else {
			l := len(br.rest)
			copy(buf, br.rest)

			br.rest = nil
			return l, nil
		}
	}

	return br.c.Read(buf)
}

func (br *unsizedBodyReader) Close() error {
	return br.c.Close()
}

func BodyReader(size int64, rest []byte, c io.ReadCloser) io.ReadCloser {
	switch size {
	case 0:
		return nil
	case -1:
		return &unsizedBodyReader{rest, c}
	default:
		return &sizedBodyReader{size, rest, c}
	}
}