/
writer.go
109 lines (98 loc) · 2.44 KB
/
writer.go
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
109
// Copyright 2015 David du Colombier. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package pcap
import (
"encoding/binary"
"errors"
"io"
)
var (
ErrWriteTooLong = errors.New("pcap: write too long")
ErrWriteAfterClose = errors.New("pcap: write after close")
)
// A Writer provides sequential writing of a pcap file.
// A pcap file consists of a sequence of records.
// Call WriteRecordHeader to begin a new record, and then call Write to supply that record's data,
// writing at most hdr.Size bytes in total.
type Writer struct {
w io.Writer
err error
nb uint32 // number of unwritten bytes for current record
closed bool
}
// NewWriter creates a new Writer writing to w.
func NewWriter(w io.Writer) *Writer { return &Writer{w: w} }
// WriteHeader writes hdr and prepares to accept the record's contents.
// Calling after a Close will return ErrWriteAfterClose.
func (pw *Writer) WriteHeader(hdr *Header) error {
hdr.Magic = Magic
hdr.VersionMajor = 2
hdr.VersionMinor = 4
if hdr.SnapLen == 0 {
hdr.SnapLen = MaxSnapLen
}
return pw.writeHeader(hdr)
}
func (pw *Writer) writeHeader(hdr *Header) error {
if pw.closed {
return ErrWriteAfterClose
}
if pw.err != nil {
return pw.err
}
err := binary.Write(pw.w, binary.LittleEndian, hdr)
if err != nil {
return err
}
return nil
}
// WriteRecordHeader writes hdr and prepares to accept the record's contents.
// Calling after a Close will return ErrWriteAfterClose.
func (pw *Writer) WriteRecordHeader(hdr *RecordHeader) error {
if pw.closed {
return ErrWriteAfterClose
}
if pw.err != nil {
return pw.err
}
err := binary.Write(pw.w, binary.LittleEndian, hdr)
if err != nil {
return err
}
pw.nb = hdr.CapLen
return nil
}
// Write writes to the current record in the pcap file.
// Write returns the error ErrWriteTooLong if more than
// hdr.CapLen bytes are written after WriteHeader.
func (pw *Writer) Write(b []byte) (n int, err error) {
if pw.closed {
err = ErrWriteAfterClose
return
}
overwrite := false
if uint32(len(b)) > pw.nb {
b = b[0:pw.nb]
overwrite = true
}
n, err = pw.w.Write(b)
pw.nb -= uint32(n)
if err == nil && overwrite {
err = ErrWriteTooLong
return
}
pw.err = err
return
}
// Close closes the pcap file.
func (pw *Writer) Close() error {
if pw.err != nil || pw.closed {
return pw.err
}
pw.closed = true
if pw.err != nil {
return pw.err
}
return pw.err
}