From 09e2c01560df5aaaed50e48e77547858e2623498 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Wed, 22 Sep 2021 21:38:58 +0800
Subject: [PATCH] fix shaper prio overflow

---
 session.go     | 4 ++--
 shaper.go      | 6 +++++-
 shaper_test.go | 6 ++++--
 stream.go      | 4 ++--
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/session.go b/session.go
index bc56066..de0c69b 100644
--- a/session.go
+++ b/session.go
@@ -24,7 +24,7 @@ var (
 )
 
 type writeRequest struct {
-	prio   uint64
+	prio   uint32
 	frame  Frame
 	result chan writeResult
 }
@@ -496,7 +496,7 @@ func (s *Session) writeFrame(f Frame) (n int, err error) {
 }
 
 // internal writeFrame version to support deadline used in keepalive
-func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint64) (int, error) {
+func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint32) (int, error) {
 	req := writeRequest{
 		prio:   prio,
 		frame:  f,
diff --git a/shaper.go b/shaper.go
index be03406..c074186 100644
--- a/shaper.go
+++ b/shaper.go
@@ -1,9 +1,13 @@
 package smux
 
+func _itimediff(later, earlier uint32) int32 {
+	return (int32)(later - earlier)
+}
+
 type shaperHeap []writeRequest
 
 func (h shaperHeap) Len() int            { return len(h) }
-func (h shaperHeap) Less(i, j int) bool  { return h[i].prio < h[j].prio }
+func (h shaperHeap) Less(i, j int) bool  { return _itimediff(h[j].prio, h[i].prio) > 0 }
 func (h shaperHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
 func (h *shaperHeap) Push(x interface{}) { *h = append(*h, x.(writeRequest)) }
 
diff --git a/shaper_test.go b/shaper_test.go
index 6f6d265..572f4ae 100644
--- a/shaper_test.go
+++ b/shaper_test.go
@@ -10,17 +10,19 @@ func TestShaper(t *testing.T) {
 	w2 := writeRequest{prio: 10}
 	w3 := writeRequest{prio: 20}
 	w4 := writeRequest{prio: 100}
+	w5 := writeRequest{prio: (1 << 32) - 1}
 
 	var reqs shaperHeap
+	heap.Push(&reqs, w5)
 	heap.Push(&reqs, w4)
 	heap.Push(&reqs, w3)
 	heap.Push(&reqs, w2)
 	heap.Push(&reqs, w1)
 
-	var lastPrio uint64
+	var lastPrio = reqs[0].prio
 	for len(reqs) > 0 {
 		w := heap.Pop(&reqs).(writeRequest)
-		if w.prio < lastPrio {
+		if int32(w.prio-lastPrio) < 0 {
 			t.Fatal("incorrect shaper priority")
 		}
 
diff --git a/stream.go b/stream.go
index 6c9499c..a95fb56 100644
--- a/stream.go
+++ b/stream.go
@@ -325,7 +325,7 @@ func (s *Stream) Write(b []byte) (n int, err error) {
 		}
 		frame.data = bts[:sz]
 		bts = bts[sz:]
-		n, err := s.sess.writeFrameInternal(frame, deadline, uint64(s.numWritten))
+		n, err := s.sess.writeFrameInternal(frame, deadline, s.numWritten)
 		s.numWritten++
 		sent += n
 		if err != nil {
@@ -393,7 +393,7 @@ func (s *Stream) writeV2(b []byte) (n int, err error) {
 				}
 				frame.data = bts[:sz]
 				bts = bts[sz:]
-				n, err := s.sess.writeFrameInternal(frame, deadline, uint64(atomic.LoadUint32(&s.numWritten)))
+				n, err := s.sess.writeFrameInternal(frame, deadline, atomic.LoadUint32(&s.numWritten))
 				atomic.AddUint32(&s.numWritten, uint32(sz))
 				sent += n
 				if err != nil {