handshack verification

handshack.go

@@ -1,98 +1,122 @@
-package sharpshooter
-
-import (
-	"github.com/soyum2222/sharpshooter/protocol"
-	"net"
-	"time"
-)
-
-func firstHandShack(h *headquarters, remote net.Addr) {
-
-	_, ok := h.Snipers.Load(remote.String())
-	if ok {
-		return
-	}
-
-	sn := NewSniper(h.conn, remote.(*net.UDPAddr))
-
-	h.Snipers.Store(remote.String(), sn)
-
-	ammo := protocol.Ammo{
-		Id:   0,
-		Kind: protocol.SECONDHANDSHACK,
-		Body: nil,
-	}
-
-	go func() {
-
-		var try int
-		ticker := time.NewTicker(time.Millisecond * 200)
-		defer ticker.Stop()
-
-	loop:
-
-		if try > DEFAULT_INIT_HANDSHACK_TIMEOUT {
-			return
-		}
-
-		sn.status = STATUS_SECONDHANDSHACK
-		_, _ = sn.conn.WriteToUDP(protocol.Marshal(ammo), sn.aim)
-
-		select {
-		case <-ticker.C:
-			try++
-			goto loop
-
-		case <-sn.handShakeSign:
-			return
-		}
-	}()
-}
-
-func secondHandShack(h *headquarters, remote net.Addr) {
-	i, ok := h.Snipers.Load(remote.String())
-	if !ok {
-		return
-	}
-
-	sn := i.(*Sniper)
-
-	// verification status
-	//if sn.status != STATUS_FIRSTHANDSHACK {
-	//	return
-	//}
-
-	sn.status = STATUS_THIRDHANDSHACK
-
-	ammo := protocol.Ammo{
-		Id:   0,
-		Kind: protocol.THIRDHANDSHACK,
-		Body: nil,
-	}
-
-	_, _ = sn.conn.WriteToUDP(protocol.Marshal(ammo), sn.aim)
-}
-
-func thirdHandShack(h *headquarters, remote net.Addr) {
-	i, ok := h.Snipers.Load(remote.String())
-	if !ok {
-		return
-	}
-
-	sn := i.(*Sniper)
-
-	//if sn.status != STATUS_SECONDHANDSHACK {
-	//	return
-	//}
-
-	sn.status = STATUS_NORMAL
-
-	SystemTimedSched.Put(sn.healthMonitor, time.Now().Add(time.Second*3))
-
-	select {
-	case sn.handShakeSign <- struct{}{}:
-	default:
-	}
-
-	h.accept <- sn
-}
+package sharpshooter
+
+import (
+	"github.com/soyum2222/sharpshooter/protocol"
+	"math/rand"
+	"net"
+	"time"
+)
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+func firstHandShack(h *headquarters, remote net.Addr) {
+
+	_, ok := h.Snipers.Load(remote.String())
+	if ok {
+		return
+	}
+
+	sn := NewSniper(h.conn, remote.(*net.UDPAddr))
+
+	// verification status
+	// avoid malicious packages
+	if sn.status != STATUS_NONE {
+		return
+	}
+
+	h.Snipers.Store(remote.String(), sn)
+
+	sn.latestSendId = rand.Uint32()
+	ammo := protocol.Ammo{
+		Id:   sn.latestSendId,
+		Kind: protocol.SECONDHANDSHACK,
+		Body: nil,
+	}
+
+	go func() {
+
+		var try int
+		ticker := time.NewTicker(time.Millisecond * 200)
+		defer ticker.Stop()
+
+	loop:
+
+		if try > DEFAULT_INIT_HANDSHACK_TIMEOUT {
+			return
+		}
+
+		sn.status = STATUS_SECONDHANDSHACK
+		_, _ = sn.conn.WriteToUDP(protocol.Marshal(ammo), sn.aim)
+
+		select {
+		case <-ticker.C:
+			try++
+			goto loop
+
+		case <-sn.handShakeSign:
+			return
+		}
+	}()
+}
+
+func secondHandShack(h *headquarters, remote net.Addr, id uint32) {
+	i, ok := h.Snipers.Load(remote.String())
+	if !ok {
+		return
+	}
+
+	sn := i.(*Sniper)
+
+	// verification status
+	// avoid malicious packages
+	if sn.status != STATUS_NONE {
+		return
+	}
+
+	sn.status = STATUS_THIRDHANDSHACK
+
+	ammo := protocol.Ammo{
+		Id:   id,
+		Kind: protocol.THIRDHANDSHACK,
+		Body: nil,
+	}
+
+	_, _ = sn.conn.WriteToUDP(protocol.Marshal(ammo), sn.aim)
+}
+
+func thirdHandShack(h *headquarters, remote net.Addr, id uint32) {
+	i, ok := h.Snipers.Load(remote.String())
+	if !ok {
+		return
+	}
+
+	sn := i.(*Sniper)
+
+	// verification status
+	// avoid malicious packages
+	if sn.status != STATUS_SECONDHANDSHACK {
+		return
+	}
+
+	if id != sn.latestSendId {
+		h.Snipers.Delete(remote.String())
+		return
+	}
+
+	sn.latestSendId = 0
+	//if sn.status != STATUS_SECONDHANDSHACK {
+	//	return
+	//}
+
+	sn.status = STATUS_NORMAL
+
+	SystemTimedSched.Put(sn.healthMonitor, time.Now().Add(time.Second*3))
+
+	select {
+	case sn.handShakeSign <- struct{}{}:
+	default:
+	}
+
+	h.accept <- sn
+}

headquarters.go

@@ -45,6 +45,8 @@ func Dial(addr string) (net.Conn, error) {
 
 	c := make(chan error, 1)
 
+	var id uint32
+
 	go func() {
 
 		secondhand := make([]byte, 14)
@@ -66,6 +68,7 @@ func Dial(addr string) (net.Conn, error) {
 			return
 		}
 
+		id = ammo.Id
 		c <- nil
 	}()
 
@@ -99,6 +102,7 @@ loop:
 	}
 
 	ammo.Kind = protocol.THIRDHANDSHACK
+	ammo.Id = id
 
 	_, err = sn.conn.WriteToUDP(protocol.Marshal(ammo), sn.aim)
 	if err != nil {
@@ -229,10 +233,10 @@ func (h *headquarters) monitor() {
 				firstHandShack(h, remote)
 
 			case protocol.SECONDHANDSHACK:
-				secondHandShack(h, remote)
+				secondHandShack(h, remote, msg.Id)
 
 			case protocol.THIRDHANDSHACK:
-				thirdHandShack(h, remote)
+				thirdHandShack(h, remote, msg.Id)
 
 			default:
 				sn, ok := h.Snipers.Load(remote.String())