99 "math/rand"
1010 "net"
1111 "sync"
12+ "sync/atomic"
1213 "time"
14+ "unsafe"
1315)
1416
1517// An ICMPRequest represents an ICMPRequest issued by ping or trace for listener
@@ -34,16 +36,20 @@ func (r *ICMPRequest) SetTimeout(duration time.Duration) {
3436 r .Deadline = r .IssueTime .Add (duration )
3537}
3638
37- func (r ICMPRequest ) Passed (time time.Time ) bool {
39+ func (r * ICMPRequest ) Passed (time time.Time ) bool {
3840 return r .Deadline .Before (time )
3941}
4042
41- func (r ICMPRequest ) Deliver (response Response ) bool {
43+ func (r * ICMPRequest ) Deliver (response Response ) bool {
44+ if r .delivery == nil {
45+ return false
46+ }
47+
4248 if response == nil {
4349 r .delivery <- & Result {
4450 Code : 256 ,
4551 }
46- close ( r .delivery )
52+ r .delivery = nil
4753 return true
4854 }
4955 ID , TargetIP := response .GetIdentifier ()
@@ -60,7 +66,7 @@ func (r ICMPRequest) Deliver(response Response) bool {
6066 Code : Code ,
6167 }
6268 }
63- close ( r .delivery )
69+ r .delivery = nil
6470 return true
6571 }
6672 return false
@@ -82,11 +88,11 @@ type ICMPResponse struct {
8288 Code int
8389}
8490
85- func (I ICMPResponse ) GetIdentifier () (int , net.IP ) {
91+ func (I * ICMPResponse ) GetIdentifier () (int , net.IP ) {
8692 return I .ID , I .TargetIP
8793}
8894
89- func (I ICMPResponse ) GetInformation () (net.IP , time.Time , int ) {
95+ func (I * ICMPResponse ) GetInformation () (net.IP , time.Time , int ) {
9096 return I .AddrIP , I .Received , I .Code
9197}
9298
@@ -118,12 +124,10 @@ type ICMPManager struct {
118124 // which send other Protocol message(e.g. TCP, UDP) but expect ICMP reply
119125 // messages.
120126 extListener map [int ]chan * RawResponse
121- // counter will fill the sequence field of the request (precisely 16bits)
127+ // counter will fill the sequence field of the request (use low 16bits)
122128 // to identify packet. it will be increased for each call and can hold at
123129 // most 65536 concurrent pending requests.
124- counter uint16
125- // l is the mutex to make counter increment thread safe.
126- l sync.Mutex
130+ counter uint32
127131 // context to send the manager stop message
128132 ctx context.Context
129133 // function to call to stop the manager
@@ -138,7 +142,7 @@ type ICMPManager struct {
138142var manager * ICMPManager
139143var once sync.Once
140144
141- // listen to ICMP socket to receive packet
145+ // ICMPv4Receiver listen to ICMP socket to receive packet
142146func ICMPv4Receiver (conn * icmp.PacketConn , wait time.Duration , icmpResponse chan * ICMPResponse ,
143147 rawResponse chan * RawResponse , ctx context.Context ) {
144148 select {
@@ -151,7 +155,11 @@ func ICMPv4Receiver(conn *icmp.PacketConn, wait time.Duration, icmpResponse chan
151155 // return
152156 //}
153157 // wait `wait` to receive some body
154- if err := conn .SetDeadline (time .Now ().Add (wait )); err != nil {
158+ var deadline time.Time
159+ if wait != 0 {
160+ deadline = time .Now ().Add (wait )
161+ }
162+ if err := conn .SetDeadline (deadline ); err != nil {
155163 return
156164 }
157165 readBytes := make ([]byte , 1500 ) // max MTU
@@ -238,7 +246,7 @@ func ICMPv4Receiver(conn *icmp.PacketConn, wait time.Duration, icmpResponse chan
238246 }
239247}
240248
241- // listen to ICMPv6 socket to receive packet
249+ // ICMPv6Receiver listen to ICMPv6 socket to receive packet
242250func ICMPv6Receiver (conn * icmp.PacketConn , wait time.Duration , icmpResponse chan * ICMPResponse ,
243251 rawResponse chan * RawResponse , ctx context.Context ) {
244252 select {
@@ -251,7 +259,11 @@ func ICMPv6Receiver(conn *icmp.PacketConn, wait time.Duration, icmpResponse chan
251259 // return
252260 //}
253261 // wait `wait` to receive some body
254- if err := conn .SetDeadline (time .Now ().Add (wait )); err != nil {
262+ var deadline time.Time
263+ if wait != 0 {
264+ deadline = time .Now ().Add (wait )
265+ }
266+ if err := conn .SetDeadline (deadline ); err != nil {
255267 return
256268 }
257269 readBytes := make ([]byte , 1500 ) // max MTU
@@ -338,8 +350,9 @@ func ICMPv6Receiver(conn *icmp.PacketConn, wait time.Duration, icmpResponse chan
338350 }
339351}
340352
341- // return ICMPManager to caller. As listening to ICMP will receive all ICMP
342- // packet, there will be only one manager in the whole process.
353+ // GetICMPManager return ICMPManager to caller.
354+ // As listening to ICMP will receive all ICMP packet,
355+ // there will be only one manager in the whole process.
343356func GetICMPManager () * ICMPManager {
344357 once .Do (func () {
345358 ctx , cancel := context .WithCancel (context .Background ())
@@ -363,15 +376,15 @@ func GetICMPManager() *ICMPManager {
363376 panic (fmt .Sprintf ("Can't listen to ICMPv6: %s" , err ))
364377 }
365378 manager .pConn6 = conn6
366- go ICMPv4Receiver (conn4 , 1000 * time . Millisecond , result4 , raw4 , ctx )
367- go ICMPv6Receiver (conn6 , 1000 * time . Millisecond , result6 , raw6 , ctx )
379+ go ICMPv4Receiver (conn4 , 0 , result4 , raw4 , ctx )
380+ go ICMPv6Receiver (conn6 , 0 , result6 , raw6 , ctx )
368381 go manager .icmpDispatcher (result4 , result6 )
369382 go manager .rawDispatcher (raw4 , raw6 )
370383 // warm-up
371384 addr , _ := net .ResolveIPAddr ("" , "127.0.0.1" )
372- manager .Issue (addr , 100 , time .Second )
385+ <- manager .Issue (addr , 100 , time .Second )
373386 addr , _ = net .ResolveIPAddr ("" , "::1" )
374- manager .Issue (addr , 100 , time .Second )
387+ <- manager .Issue (addr , 100 , time .Second )
375388 })
376389 return manager
377390}
@@ -389,10 +402,7 @@ func (mgr *ICMPManager) Issue(ip net.Addr, ttl int, timeout time.Duration) (deli
389402 }
390403 dest = ipAddr .IP .To16 ()
391404
392- mgr .l .Lock ()
393- count := mgr .counter
394- mgr .counter ++
395- mgr .l .Unlock ()
405+ count := (atomic .AddUint32 (& mgr .counter , 1 ) - 1 ) & 0xffff
396406
397407 id := rand .Intn (1 << 16 )
398408 var msg []byte
@@ -500,6 +510,14 @@ func (mgr *ICMPManager) rawDispatcher(v4, v6 chan *RawResponse) {
500510 }
501511}
502512
513+ func (mgr * ICMPManager ) Flush () {
514+ queue := NewCMap (32 )
515+ queue = (* ConMapRequest )(atomic .SwapPointer ((* unsafe .Pointer )((unsafe .Pointer )(& mgr .queue )), unsafe .Pointer (queue )))
516+ for t := range mgr .queue .IterBuffered () {
517+ t .Val .Deliver (nil )
518+ }
519+ }
520+
503521func (mgr * ICMPManager ) Finish () {
504522 mgr .cancel ()
505523}
0 commit comments