Merge pull request #2 from AudriusButkevicius/small-reads

Handle small reads
AudriusButkevicius · Jun 27, 2019 · c55ef61 · c55ef61
2 parents 730b0de + ff3c123
commit c55ef61
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/conn.go b/conn.go
@@ -76,9 +76,13 @@ func (r *FilteredConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
 	case <-timeout:
 		return 0, nil, errTimeout
 	case pkt := <-r.recvBuffer:
-		copy(b[:pkt.n], pkt.buf)
+		n := pkt.n
+		if l := len(b); l < n {
+			n = l
+		}
+		copy(b, pkt.buf[:n])
 		bufPool.Put(pkt.buf[:maxPacketSize])
-		return pkt.n, pkt.addr, pkt.err
+		return n, pkt.addr, pkt.err
 	case <-r.closed:
 		return 0, nil, errClosed
 	}

diff --git a/filter.go b/filter.go
@@ -124,12 +124,13 @@ next:
 				select {
 				case conn.recvBuffer <- pkt:
 				default:
+					bufPool.Put(pkt.buf[:maxPacketSize])
 					atomic.AddUint64(&d.overflow, 1)
 				}
 				goto next
 			}
 		}
-
+		bufPool.Put(pkt.buf[:maxPacketSize])
 		atomic.AddUint64(&d.dropped, 1)
 	}
 }