lzhr/v2fly
1
0
Fork 0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 18:17:52 -05:00
Code Issues Releases Wiki Activity

check connection state for every write operation

Browse Source
This commit is contained in:
Darien Raymond 2018-01-17 16:18:38 +01:00
parent b5caea67ac
commit a6c0ef11ba
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 33 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
transport/internet/kcp/connection.go
View File

@ -340,22 +340,31 @@ func (c *Connection) waitForDataOutput() error {
func (c *Connection) Write(b []byte) (int, error) { func (c *Connection) Write(b []byte) (int, error) {
totalWritten := 0 totalWritten := 0
for {
dataWritten := false
for { for {
if c == nil || c.State() != StateActive { if c == nil || c.State() != StateActive {
return totalWritten, io.ErrClosedPipe return totalWritten, io.ErrClosedPipe
} }
if !c.sendingWorker.Push(func(bb []byte) (int, error) {
for c.sendingWorker.Push(func(bb []byte) (int, error) {
n := copy(bb[:c.mss], b[totalWritten:]) n := copy(bb[:c.mss], b[totalWritten:])
totalWritten += n totalWritten += n
return n, nil return n, nil
}) { }) {
c.dataUpdater.WakeUp() break
}
dataWritten = true
if totalWritten == len(b) { if totalWritten == len(b) {
return totalWritten, nil return totalWritten, nil
} }
} }
if dataWritten {
c.dataUpdater.WakeUp()
}
if err := c.waitForDataOutput(); err != nil { if err := c.waitForDataOutput(); err != nil {
return totalWritten, err return totalWritten, err
} }
@ -366,20 +375,28 @@ func (c *Connection) Write(b []byte) (int, error) {
func (c *Connection) WriteMultiBuffer(mb buf.MultiBuffer) error { func (c *Connection) WriteMultiBuffer(mb buf.MultiBuffer) error {
defer mb.Release() defer mb.Release()
for {
dataWritten := false
for { for {
if c == nil || c.State() != StateActive { if c == nil || c.State() != StateActive {
return io.ErrClosedPipe return io.ErrClosedPipe
} }
for c.sendingWorker.Push(func(bb []byte) (int, error) { if !c.sendingWorker.Push(func(bb []byte) (int, error) {
return mb.Read(bb[:c.mss]) return mb.Read(bb[:c.mss])
}) { }) {
c.dataUpdater.WakeUp() break
}
dataWritten = true
if mb.IsEmpty() { if mb.IsEmpty() {
return nil return nil
} }
} }
if dataWritten {
c.dataUpdater.WakeUp()
}
if err := c.waitForDataOutput(); err != nil { if err := c.waitForDataOutput(); err != nil {
return err return err
} }