github.com/cnotch/ipchub@v1.1.0/network/socket/buffered/conn.go (about) 1 // Copyright (c) 2019,CAOHONGJU All rights reserved. 2 // Use of this source code is governed by a MIT-style 3 // license that can be found in the LICENSE file. 4 5 package buffered 6 7 import ( 8 "bufio" 9 "bytes" 10 "net" 11 "time" 12 13 "github.com/kelindar/rate" 14 ) 15 16 const ( 17 defaultRate = 50 18 defaultBufferSize = 64 * 1024 19 minBufferSize = 8 * 1024 20 ) 21 22 // Conn wraps a net.Conn and provides buffered ability. 23 type Conn struct { 24 socket net.Conn // The underlying network connection. 25 reader *bufio.Reader // The buffered reader 26 writer *bytes.Buffer // The buffered write queue. 27 limit *rate.Limiter // The write rate limiter. 28 bufferSize int // The read and write max buffer size 29 } 30 31 // NewConn creates a new sniffed connection. 32 func NewConn(c net.Conn, options ...Option) *Conn { 33 conn, ok := c.(*Conn) 34 if !ok { 35 conn = &Conn{ 36 socket: c, 37 } 38 } 39 40 for _, option := range options { 41 option.apply(conn) 42 } 43 44 // 设置默认值刷新频率 45 if conn.limit == nil { 46 conn.limit = rate.New(defaultRate, time.Second) 47 } 48 49 if conn.bufferSize <= 0 { 50 conn.bufferSize = defaultBufferSize 51 } 52 53 // 设置IO缓冲对象 54 conn.reader = bufio.NewReaderSize(conn.socket, conn.bufferSize) 55 conn.writer = bytes.NewBuffer(make([]byte, 0, conn.bufferSize)) 56 return conn 57 } 58 59 // Buffered returns the pending buffer size. 60 func (m *Conn) Buffered() (n int) { 61 return m.writer.Len() 62 } 63 64 // Reader 返回内部的 bufio.Reader 65 func (m *Conn) Reader() *bufio.Reader { 66 return m.reader 67 } 68 69 // Flush flushes the underlying buffer by writing into the underlying connection. 70 func (m *Conn) Flush() (n int, err error) { 71 if m.Buffered() == 0 { 72 return 0, nil 73 } 74 75 // Flush everything and reset the buffer 76 n, err = m.writeFull(m.writer.Bytes()) 77 m.writer.Reset() 78 return 79 } 80 81 // Read reads the block of data from the underlying buffer. 82 func (m *Conn) Read(p []byte) (int, error) { 83 return m.reader.Read(p) 84 } 85 86 // Write writes the block of data into the underlying buffer. 87 func (m *Conn) Write(p []byte) (nn int, err error) { 88 var n int 89 // 没有足够的空间容纳 p 90 for len(p) > m.bufferSize-m.Buffered() && err == nil { 91 if m.Buffered() == 0 { 92 // Large write, empty buffer. 93 // Write directly from p to avoid copy. 94 n, err = m.socket.Write(p) 95 } else { 96 // write buffer to full state,and flush 97 n, err = m.writer.Write(p[:m.bufferSize-m.writer.Len()]) 98 _, err = m.Flush() 99 } 100 nn += n 101 p = p[n:] 102 } 103 104 if err != nil { 105 return nn, err 106 } 107 108 // 未到达时间频率的间隔,直接写到缓存 109 if m.limit.Limit() { 110 n, err = m.writer.Write(p) 111 return nn + n, err 112 } 113 114 // 缓存中有数据,flush 115 if m.Buffered() > 0 { 116 n, err = m.writer.Write(p) 117 _, err = m.Flush() 118 return nn + n, err 119 } 120 121 // 缓存中无数据,直接写避免内存拷贝 122 n, err = m.writeFull(p) 123 return nn + n, err 124 125 } 126 127 func (m *Conn) writeFull(p []byte) (nn int, err error) { 128 var n int 129 for len(p) > 0 && err == nil { 130 n, err = m.socket.Write(p) 131 nn += n 132 p = p[n:] 133 } 134 return nn, err 135 } 136 137 // Close closes the connection. Any blocked Read or Write operations will be unblocked 138 // and return errors. 139 func (m *Conn) Close() error { 140 return m.socket.Close() 141 } 142 143 // LocalAddr returns the local network address. 144 func (m *Conn) LocalAddr() net.Addr { 145 return m.socket.LocalAddr() 146 } 147 148 // RemoteAddr returns the remote network address. 149 func (m *Conn) RemoteAddr() net.Addr { 150 return m.socket.RemoteAddr() 151 } 152 153 // SetDeadline sets the read and write deadlines associated 154 // with the connection. It is equivalent to calling both 155 // SetReadDeadline and SetWriteDeadline. 156 func (m *Conn) SetDeadline(t time.Time) error { 157 return m.socket.SetDeadline(t) 158 } 159 160 // SetReadDeadline sets the deadline for future Read calls 161 // and any currently-blocked Read call. 162 func (m *Conn) SetReadDeadline(t time.Time) error { 163 return m.socket.SetReadDeadline(t) 164 } 165 166 // SetWriteDeadline sets the deadline for future Write calls 167 // and any currently-blocked Write call. 168 func (m *Conn) SetWriteDeadline(t time.Time) error { 169 return m.socket.SetWriteDeadline(t) 170 } 171 172 // Option 配置 Conn 的选项接口 173 type Option interface { 174 apply(*Conn) 175 } 176 177 // OptionFunc 包装函数以便它满足 Option 接口 178 type optionFunc func(*Conn) 179 180 func (f optionFunc) apply(c *Conn) { 181 f(c) 182 } 183 184 // FlushRate Conn 写操作的每秒刷新频率 185 func FlushRate(r int) Option { 186 return optionFunc(func(c *Conn) { 187 if r < 1 { // 如果不合规,设置成默认值 188 r = defaultRate 189 } 190 c.limit = rate.New(r, time.Second) 191 }) 192 } 193 194 // BufferSize Conn 缓冲大小 195 func BufferSize(bufferSize int) Option { 196 return optionFunc(func(c *Conn) { 197 if bufferSize < minBufferSize { // 如果不合规,设置成最小值 198 bufferSize = minBufferSize 199 } 200 c.bufferSize = bufferSize 201 }) 202 }