From 7033f3e72b43f3364300c70449d92c5355614be3 Mon Sep 17 00:00:00 2001
From: Guy Lewin <guy.lewin@microsoft.com>
Date: Mon, 10 Feb 2020 19:14:07 +0200
Subject: [PATCH] Test TlsOnly

---
 tests/ci/tls_test.go | 92 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/tests/ci/tls_test.go b/tests/ci/tls_test.go
index d2ad801..c46e9fa 100644
--- a/tests/ci/tls_test.go
+++ b/tests/ci/tls_test.go
@@ -186,3 +186,95 @@ func TestTLSOverWebsocket(t *testing.T) {
 	assert.NoError(err)
 	assert.Equal(consts.TEST_TCP_ECHO_STR, res)
 }
+
+const FRPS_TLS_ONLY_TCP_CONF = `
+[common]
+bind_addr = 0.0.0.0
+bind_port = 20000
+log_file = console
+log_level = debug
+token = 123456
+tls_only = true
+`
+
+const FRPC_TLS_ONLY_TCP_CONF = `
+[common]
+server_addr = 127.0.0.1
+server_port = 20000
+log_file = console
+log_level = debug
+token = 123456
+protocol = tcp
+tls_enable = true
+
+[tcp]
+type = tcp
+local_port = 10701
+remote_port = 20801
+`
+
+const FRPC_TLS_ONLY_NO_TLS_TCP_CONF = `
+[common]
+server_addr = 127.0.0.1
+server_port = 20000
+log_file = console
+log_level = debug
+token = 123456
+protocol = tcp
+tls_enable = false
+
+[tcp]
+type = tcp
+local_port = 10701
+remote_port = 20802
+`
+
+func TestTlsOnlyOverTCP(t *testing.T) {
+	assert := assert.New(t)
+	frpsCfgPath, err := config.GenerateConfigFile(consts.FRPS_NORMAL_CONFIG, FRPS_TLS_ONLY_TCP_CONF)
+	if assert.NoError(err) {
+		defer os.Remove(frpsCfgPath)
+	}
+
+	frpcWithTlsCfgPath, err := config.GenerateConfigFile(consts.FRPC_NORMAL_CONFIG, FRPC_TLS_ONLY_TCP_CONF)
+	if assert.NoError(err) {
+		defer os.Remove(frpcWithTlsCfgPath)
+	}
+
+	frpsProcess := util.NewProcess(consts.FRPS_BIN_PATH, []string{"-c", frpsCfgPath})
+	err = frpsProcess.Start()
+	if assert.NoError(err) {
+		defer frpsProcess.Stop()
+	}
+
+	time.Sleep(200 * time.Millisecond)
+
+	frpcProcessWithTls := util.NewProcess(consts.FRPC_BIN_PATH, []string{"-c", frpcWithTlsCfgPath})
+	err = frpcProcessWithTls.Start()
+	if assert.NoError(err) {
+		defer frpcProcessWithTls.Stop()
+	}
+	time.Sleep(500 * time.Millisecond)
+
+	// test tcp over tls
+	res, err := util.SendTcpMsg("127.0.0.1:20801", consts.TEST_TCP_ECHO_STR)
+	assert.NoError(err)
+	assert.Equal(consts.TEST_TCP_ECHO_STR, res)
+	frpcProcessWithTls.Stop()
+
+	frpcWithoutTlsCfgPath, err := config.GenerateConfigFile(consts.FRPC_NORMAL_CONFIG, FRPC_TLS_ONLY_NO_TLS_TCP_CONF)
+	if assert.NoError(err) {
+		defer os.Remove(frpcWithTlsCfgPath)
+	}
+
+	frpcProcessWithoutTls := util.NewProcess(consts.FRPC_BIN_PATH, []string{"-c", frpcWithoutTlsCfgPath})
+	err = frpcProcessWithoutTls.Start()
+	if assert.NoError(err) {
+		defer frpcProcessWithoutTls.Stop()
+	}
+	time.Sleep(500 * time.Millisecond)
+
+	// test tcp without tls
+	_, err = util.SendTcpMsg("127.0.0.1:20802", consts.TEST_TCP_ECHO_STR)
+	assert.Error(err)
+}