fix shaper prio overflow

Author: xtaci

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,

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)) }
 

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")
 		}
 

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 {