From 6fe7463bf4529f0085f50d4d9d6f27cbdcc8338c Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Tue, 26 Jul 2016 21:34:00 +0200
Subject: [PATCH] simplify state checking

---
 transport/internet/kcp/connection.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/transport/internet/kcp/connection.go b/transport/internet/kcp/connection.go
index 6d91337d..d1cc34ce 100644
--- a/transport/internet/kcp/connection.go
+++ b/transport/internet/kcp/connection.go
@@ -21,6 +21,15 @@ var (
 
 type State int32
 
+func (this State) Is(states ...State) bool {
+	for _, state := range states {
+		if this == state {
+			return true
+		}
+	}
+	return false
+}
+
 const (
 	StateActive          State = 0
 	StateReadyToClose    State = 1
@@ -172,7 +181,7 @@ func (this *Connection) Read(b []byte) (int, error) {
 	}
 
 	for {
-		if this.State() == StateReadyToClose || this.State() == StateTerminating || this.State() == StateTerminated {
+		if this.State().Is(StateReadyToClose, StateTerminating, StateTerminated) {
 			return 0, io.EOF
 		}
 		nBytes := this.receivingWorker.Read(b)
@@ -206,9 +215,6 @@ func (this *Connection) Read(b []byte) (int, error) {
 
 // Write implements the Conn Write method.
 func (this *Connection) Write(b []byte) (int, error) {
-	if this == nil || this.State() != StateActive {
-		return 0, io.ErrClosedPipe
-	}
 	totalWritten := 0
 
 	for {
@@ -278,9 +284,7 @@ func (this *Connection) Close() error {
 	this.dataOutputCond.Broadcast()
 
 	state := this.State()
-	if state == StateReadyToClose ||
-		state == StateTerminating ||
-		state == StateTerminated {
+	if state.Is(StateReadyToClose, StateTerminating, StateTerminated) {
 		return errClosedConnection
 	}
 	log.Info("KCP|Connection: Closing connection to ", this.remote)