add 'root', 'index' config, allow tunnel read from local static files

TeaWeb · Jul 7, 2019 · c46bd75 · c46bd75
1 parent 18ec804
commit c46bd75
Show file tree

Hide file tree

Showing 4 changed files with 255 additions and 43 deletions.
diff --git a/tests/index.html b/tests/index.html
@@ -0,0 +1,9 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+	<title>Hello</title>
+</head>
+<body>
+<strong>this is html message</strong>
+</body>
+</html>
diff --git a/tunnel.go b/tunnel.go
@@ -4,14 +4,21 @@ import (
 	"bufio"
 	"bytes"
 	"errors"
+	"fmt"
+	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/logs"
+	stringutil "github.com/iwind/TeaGo/utils/string"
 	"io"
 	"io/ioutil"
 	"log"
+	"mime"
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"os"
+	"path/filepath"
 	"runtime"
+	"strings"
 	"sync"
 	"time"
 )
@@ -31,13 +38,17 @@ func NewTunnel(config *TunnelConfig) *Tunnel {
 }
 
 func (this *Tunnel) Start() error {
-	host := this.config.LocalHost()
+	localHost := this.config.LocalHost()
+	root := this.config.Root
 	scheme := this.config.LocalScheme()
 
-	if len(host) == 0 {
-		return errors.New("local host should not be empty")
+	if len(localHost) == 0 && len(root) == 0 {
+		return errors.New("'local' or 'root' should not be empty")
 	}
 
+	hasLocal := len(localHost) > 0
+	hasRoot := len(root) > 0
+
 	if len(scheme) == 0 {
 		scheme = "http"
 	}
@@ -114,55 +125,215 @@ func (this *Tunnel) Start() error {
 					}
 				}
 
-				req.RequestURI = ""
-				req.URL.Host = host
-				req.URL.Scheme = scheme
-
 				logs.Println(req.Header.Get("X-Forwarded-For") + " - \"" + req.Method + " " + req.URL.String() + "\" \"" + req.Header.Get("User-Agent") + "\"")
 
-				if len(this.config.Host) > 0 {
-					req.Host = this.config.Host
-				} else {
-					forwardedHost := req.Header.Get("X-Forwarded-Host")
-					if len(forwardedHost) > 0 {
-						req.Host = forwardedHost
+				if hasLocal { // read from local web server
+					req.RequestURI = ""
+					req.URL.Host = localHost
+					req.URL.Scheme = scheme
+
+					if len(this.config.Host) > 0 {
+						req.Host = this.config.Host
 					} else {
-						req.Host = host
+						forwardedHost := req.Header.Get("X-Forwarded-Host")
+						if len(forwardedHost) > 0 {
+							req.Host = forwardedHost
+						} else {
+							req.Host = localHost
+						}
 					}
-				}
 
-				resp, err := HttpClient.Do(req)
-				if err != nil {
-					logs.Error(err)
-					resp := &http.Response{
-						StatusCode: http.StatusBadGateway,
-						Status:     "Bad Gateway",
-						Header: map[string][]string{
-							"Content-Type": {"text/plain"},
-							"Connection":   {"keep-alive"},
-						},
-						Proto:      "HTTP/1.1",
-						ProtoMajor: 1,
-						ProtoMinor: 1,
-					}
-					data, err := httputil.DumpResponse(resp, false)
+					resp, err := HttpClient.Do(req)
 					if err != nil {
 						logs.Error(err)
+						resp := &http.Response{
+							StatusCode: http.StatusBadGateway,
+							Status:     "Bad Gateway",
+							Header: map[string][]string{
+								"Content-Type": {"text/plain"},
+								"Connection":   {"keep-alive"},
+							},
+							Proto:      "HTTP/1.1",
+							ProtoMajor: 1,
+							ProtoMinor: 1,
+						}
+						data, err := httputil.DumpResponse(resp, false)
+						if err != nil {
+							logs.Error(err)
+							this.writeServerError(conn)
+							continue
+						}
+						conn.Write(data)
+					} else {
+						resp.Header.Set("Connection", "keep-alive")
+						data, err := httputil.DumpResponse(resp, true)
+						if err != nil {
+							logs.Error(err)
+							resp.Body.Close()
+							this.writeServerError(conn)
+							continue
+						}
+						conn.Write(data)
+						resp.Body.Close()
+					}
+				} else if hasRoot { // read from root directory
+					requestPath := req.URL.Path
+					if len(requestPath) == 0 || requestPath == "/" {
+						path, stat, found := this.findIndexPage(root + Tea.DS)
+						if found {
+							this.writeFile(conn, req, stat, path)
+							continue
+						}
+						this.writeNotFound(conn)
+						continue
+					}
+
+					path := root + Tea.DS + requestPath
+
+					if strings.HasSuffix(path, "/") {
+						path, stat, found := this.findIndexPage(path)
+						if found {
+							this.writeFile(conn, req, stat, path)
+							continue
+						}
+						this.writeNotFound(conn)
 						continue
 					}
-					conn.Write(data)
-				} else {
-					resp.Header.Set("Connection", "keep-alive")
-					data, err := httputil.DumpResponse(resp, true)
+
+					stat, err := os.Stat(path)
 					if err != nil {
-						logs.Error(err)
+						if os.IsNotExist(err) {
+							this.writeString(conn, http.StatusNotFound, "File Not Found")
+						} else {
+							logs.Error(err)
+							this.writeServerError(conn)
+						}
+						continue
+					} else if stat.IsDir() { // try again
+						path, stat, found := this.findIndexPage(path)
+						if found {
+							this.writeFile(conn, req, stat, path)
+							continue
+						}
+						this.writeNotFound(conn)
 						continue
 					}
-					conn.Write(data)
-					resp.Body.Close()
+
+					this.writeFile(conn, req, stat, path)
 				}
 			}
 		}(conn)
 	}
 	return nil
 }
+
+func (this *Tunnel) writeString(conn net.Conn, code int, data string) {
+	dataBytes := []byte(data)
+	this.writeBytes(conn, code, dataBytes)
+}
+
+func (this *Tunnel) writeBytes(conn net.Conn, code int, data []byte) {
+	resp := &http.Response{
+		StatusCode:    code,
+		ContentLength: int64(len(data)),
+		ProtoMajor:    1,
+		ProtoMinor:    1,
+		Header: map[string][]string{
+			"Content-Type": {"text/html; charset=utf-8"},
+		},
+	}
+	resp.Body = ioutil.NopCloser(bytes.NewReader(data))
+	respData, err := httputil.DumpResponse(resp, true)
+	if err != nil {
+		logs.Error(err)
+		resp.Body.Close()
+		return
+	}
+
+	_, err = conn.Write(respData)
+	if err != nil {
+		logs.Error(err)
+	}
+	resp.Body.Close()
+}
+
+func (this *Tunnel) writeServerError(conn net.Conn) {
+	this.writeString(conn, http.StatusInternalServerError, "Internal Server Error")
+}
+
+func (this *Tunnel) writeNotFound(conn net.Conn) {
+	this.writeString(conn, http.StatusNotFound, "File Not Found")
+}
+
+func (this *Tunnel) writeFile(conn net.Conn, req *http.Request, stat os.FileInfo, path string) {
+	reader, err := os.OpenFile(path, os.O_RDONLY, 444)
+	if err != nil {
+		logs.Error(err)
+		this.writeServerError(conn)
+		return
+	}
+	defer reader.Close()
+
+	resp := &http.Response{
+		StatusCode:    200,
+		ContentLength: stat.Size(),
+		ProtoMajor:    1,
+		ProtoMinor:    1,
+		Header:        map[string][]string{},
+	}
+
+	// mime type
+	ext := filepath.Ext(path)
+	if len(ext) > 0 {
+		mimeType := mime.TypeByExtension(ext)
+		if len(mimeType) > 0 {
+			resp.Header.Set("Content-Type", mimeType)
+		}
+	}
+
+	// supports Last-Modified
+	modifiedTime := stat.ModTime().Format("Mon, 02 Jan 2006 15:04:05 GMT")
+	resp.Header.Set("Last-Modified", modifiedTime)
+
+	// supports ETag
+	eTag := "\"et" + stringutil.Md5(fmt.Sprintf("%d,%d", stat.ModTime().UnixNano(), stat.Size())) + "\""
+	resp.Header.Set("ETag", eTag)
+
+	// supports If-None-Match
+	if req.Header.Get("If-None-Match") == eTag {
+		this.writeBytes(conn, http.StatusNotModified, []byte{})
+		return
+	}
+
+	// supports If-Modified-Since
+	if req.Header.Get("If-Modified-Since") == modifiedTime {
+		this.writeBytes(conn, http.StatusNotModified, []byte{})
+		return
+	}
+
+	//write body
+	resp.Body = reader
+	data, err := httputil.DumpResponse(resp, true)
+	if err != nil {
+		logs.Error(err)
+		this.writeServerError(conn)
+		return
+	}
+
+	conn.Write(data)
+}
+
+func (this *Tunnel) findIndexPage(dir string) (path string, stat os.FileInfo, found bool) {
+	if len(this.config.Index) == 0 {
+		this.config.Index = []string{"index.html", "default.html"}
+	}
+	for _, index := range this.config.Index {
+		path := dir + Tea.DS + index
+		stat, err := os.Stat(path)
+		if err != nil || stat.IsDir() {
+			continue
+		}
+		return path, stat, true
+	}
+	return
+}
diff --git a/tunnel_config.go b/tunnel_config.go
@@ -1,18 +1,29 @@
 package tunnel_client
 
-import "net/url"
+import (
+	"errors"
+	"net/url"
+)
 
 type TunnelConfig struct {
-	Remote string `yaml:"remote"`
-	Local  string `yaml:"local"`
-	Host   string `yaml:"host"`
-	Secret string `yaml:"secret"`
+	Remote string   `yaml:"remote"`
+	Local  string   `yaml:"local"`
+	Host   string   `yaml:"host"`
+	Secret string   `yaml:"secret"`
+	Root   string   `yaml:"root"`  // static files root
+	Index  []string `yaml:"index"` // default index page filenames
 
 	localHost   string
 	localScheme string
 }
 
 func (this *TunnelConfig) Validate() error {
+	if len(this.Remote) == 0 {
+		return errors.New("'remote' should not be empty")
+	}
+	if len(this.Local) == 0 && len(this.Root) == 0 {
+		return errors.New("'local' or 'root' should not be empty")
+	}
 	u, err := url.Parse(this.Local)
 	if err != nil {
 		return err

diff --git a/tunnel_test.go b/tunnel_test.go
@@ -1,6 +1,8 @@
 package tunnel_client
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestTunnel_Start(t *testing.T) {
 	config := &TunnelConfig{
@@ -20,3 +22,22 @@ func TestTunnel_Start(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestTunnel_StartRoot(t *testing.T) {
+	config := &TunnelConfig{
+		Remote: "192.168.2.40:8884",
+		Host:   "www.teaos.cn",
+		Root:   ".",
+		Secret: "YKCXgsGlDcZv7o5VEjF2iT5K4t3ae5bE",
+	}
+	err := config.Validate()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tunnel := NewTunnel(config)
+	err = tunnel.Start()
+	if err != nil {
+		t.Fatal(err)
+	}
+}