From 86d7d510231f16cd1f7fb18536aab294c0717b58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 21 Sep 2022 18:16:04 +0800
Subject: [PATCH] Add force_stdio build flag

---
 common/bufio/conn.go            | 26 +++++++++++++++++++-------
 common/bufio/conn_stdio.go      |  6 ++++++
 common/bufio/conn_stdio_stub.go |  5 +++++
 3 files changed, 30 insertions(+), 7 deletions(-)
 create mode 100644 common/bufio/conn_stdio.go
 create mode 100644 common/bufio/conn_stdio_stub.go

diff --git a/common/bufio/conn.go b/common/bufio/conn.go
index 59a2174..bde9bd5 100644
--- a/common/bufio/conn.go
+++ b/common/bufio/conn.go
@@ -11,10 +11,10 @@ import (
 )
 
 func NewPacketConn(conn net.PacketConn) N.NetPacketConn {
-	if packetConn, ok := conn.(N.NetPacketConn); ok {
-		return packetConn
-	} else if udpConn, ok := conn.(*net.UDPConn); ok {
+	if udpConn, isUDPConn := conn.(*net.UDPConn); isUDPConn {
 		return &ExtendedUDPConn{udpConn}
+	} else if packetConn, isPacketConn := conn.(N.NetPacketConn); isPacketConn && !forceSTDIO {
+		return packetConn
 	} else {
 		return &ExtendedPacketConn{conn}
 	}
@@ -159,8 +159,14 @@ func (r *ExtendedReaderWrapper) ReaderReplaceable() bool {
 }
 
 func NewExtendedReader(reader io.Reader) N.ExtendedReader {
-	if r, ok := reader.(N.ExtendedReader); ok {
-		return r
+	if forceSTDIO {
+		if r, ok := reader.(*ExtendedReaderWrapper); ok {
+			return r
+		}
+	} else {
+		if r, ok := reader.(N.ExtendedReader); ok {
+			return r
+		}
 	}
 	return &ExtendedReaderWrapper{reader}
 }
@@ -187,8 +193,14 @@ func (w *ExtendedReaderWrapper) WriterReplaceable() bool {
 }
 
 func NewExtendedWriter(writer io.Writer) N.ExtendedWriter {
-	if w, ok := writer.(N.ExtendedWriter); ok {
-		return w
+	if forceSTDIO {
+		if w, ok := writer.(*ExtendedWriterWrapper); ok {
+			return w
+		}
+	} else {
+		if w, ok := writer.(N.ExtendedWriter); ok {
+			return w
+		}
 	}
 	return &ExtendedWriterWrapper{writer}
 }
diff --git a/common/bufio/conn_stdio.go b/common/bufio/conn_stdio.go
new file mode 100644
index 0000000..07985d1
--- /dev/null
+++ b/common/bufio/conn_stdio.go
@@ -0,0 +1,6 @@
+//go:build force_stdio
+
+package bufio
+
+// force_stdio is dedicated to testing exposed stdio APIs.
+const forceSTDIO = true
diff --git a/common/bufio/conn_stdio_stub.go b/common/bufio/conn_stdio_stub.go
new file mode 100644
index 0000000..0525b96
--- /dev/null
+++ b/common/bufio/conn_stdio_stub.go
@@ -0,0 +1,5 @@
+//go:build !force_stdio
+
+package bufio
+
+const forceSTDIO = false