Skip to content

Commit

Permalink
fix shaper prio overflow
Browse files Browse the repository at this point in the history
  • Loading branch information
xtaci committed Sep 22, 2021
1 parent eba6ee1 commit 09e2c01
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 7 deletions.
4 changes: 2 additions & 2 deletions session.go
Expand Up @@ -24,7 +24,7 @@ var (
)

type writeRequest struct {
prio uint64
prio uint32
frame Frame
result chan writeResult
}
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion 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)) }

Expand Down
6 changes: 4 additions & 2 deletions shaper_test.go
Expand Up @@ -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")
}

Expand Down
4 changes: 2 additions & 2 deletions stream.go
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 09e2c01

Please sign in to comment.