From 8107a8bd1e3cc71504d2bc08b03bf5af29b91ac3 Mon Sep 17 00:00:00 2001
From: stffabi <stffabi@users.noreply.github.com>
Date: Tue, 30 Nov 2021 11:10:30 +0100
Subject: [PATCH] [v2] Improve translation of URI to file to be loaded

---
 .../frontend/desktop/common/uri_translate.go  | 20 +++++++++++++++++++
 .../frontend/desktop/darwin/frontend.go       | 16 ++++++++++-----
 .../frontend/desktop/linux/frontend.go        | 12 ++++++-----
 .../frontend/desktop/windows/frontend.go      | 11 ++++++----
 4 files changed, 45 insertions(+), 14 deletions(-)
 create mode 100644 v2/internal/frontend/desktop/common/uri_translate.go

diff --git a/v2/internal/frontend/desktop/common/uri_translate.go b/v2/internal/frontend/desktop/common/uri_translate.go
new file mode 100644
index 000000000..4cfd78824
--- /dev/null
+++ b/v2/internal/frontend/desktop/common/uri_translate.go
@@ -0,0 +1,20 @@
+package common
+
+import "net/url"
+
+func TranslateUriToFile(uri string, expectedScheme string, expectedHost string) (file string, match bool, err error) {
+	url, err := url.Parse(uri)
+	if err != nil {
+		return "", false, err
+	}
+
+	if url.Scheme != expectedScheme || url.Host != expectedHost {
+		return "", false, nil
+	}
+
+	filePath := url.Path
+	if filePath == "" {
+		filePath = "/"
+	}
+	return filePath, true, nil
+}
diff --git a/v2/internal/frontend/desktop/darwin/frontend.go b/v2/internal/frontend/desktop/darwin/frontend.go
index 6acf61160..8922cd228 100644
--- a/v2/internal/frontend/desktop/darwin/frontend.go
+++ b/v2/internal/frontend/desktop/darwin/frontend.go
@@ -19,12 +19,12 @@ import (
 	"html/template"
 	"log"
 	"strconv"
-	"strings"
 	"unsafe"
 
 	"github.com/wailsapp/wails/v2/internal/binding"
 	"github.com/wailsapp/wails/v2/internal/frontend"
 	"github.com/wailsapp/wails/v2/internal/frontend/assetserver"
+	"github.com/wailsapp/wails/v2/internal/frontend/desktop/common"
 	"github.com/wailsapp/wails/v2/internal/logger"
 	"github.com/wailsapp/wails/v2/pkg/options"
 )
@@ -270,12 +270,18 @@ func (f *Frontend) ExecJS(js string) {
 }
 
 func (f *Frontend) processRequest(r *request) {
-	url := C.GoString(r.url)
-	url = strings.TrimPrefix(url, "wails://wails")
-	if !strings.HasPrefix(url, "/") {
+	uri := C.GoString(r.url)
+
+	// Translate URI to file
+	file, match, err := common.TranslateUriToFile(uri, "wails", "wails")
+	if err != nil {
+		// TODO Handle errors
+		return
+	} else if !match {
 		return
 	}
-	_contents, _mimetype, err := f.assets.Load(url)
+
+	_contents, _mimetype, err := f.assets.Load(file)
 	if err != nil {
 		f.logger.Error(err.Error())
 		//TODO: Handle errors
diff --git a/v2/internal/frontend/desktop/linux/frontend.go b/v2/internal/frontend/desktop/linux/frontend.go
index d082a8f1d..7523396f7 100644
--- a/v2/internal/frontend/desktop/linux/frontend.go
+++ b/v2/internal/frontend/desktop/linux/frontend.go
@@ -27,7 +27,6 @@ import (
 	"log"
 	"os"
 	"strconv"
-	"strings"
 	"sync"
 	"text/template"
 	"unsafe"
@@ -35,6 +34,7 @@ import (
 	"github.com/wailsapp/wails/v2/internal/binding"
 	"github.com/wailsapp/wails/v2/internal/frontend"
 	"github.com/wailsapp/wails/v2/internal/frontend/assetserver"
+	"github.com/wailsapp/wails/v2/internal/frontend/desktop/common"
 	"github.com/wailsapp/wails/v2/internal/logger"
 	"github.com/wailsapp/wails/v2/pkg/options"
 )
@@ -328,14 +328,16 @@ func (f *Frontend) processRequest(request unsafe.Pointer) {
 	uri := C.webkit_uri_scheme_request_get_uri(req)
 	goURI := C.GoString(uri)
 
-	// Translate URI
-	goURI = strings.TrimPrefix(goURI, "wails://")
-	if !strings.HasPrefix(goURI, "/") {
+	file, match, err := common.TranslateUriToFile(uri, "wails", "")
+	if err != nil {
+		// TODO Handle errors
+		return
+	} else if !match {
 		return
 	}
 
 	// Load file from asset store
-	content, mimeType, err := f.assets.Load(goURI)
+	content, mimeType, err := f.assets.Load(file)
 	if err != nil {
 		return
 	}
diff --git a/v2/internal/frontend/desktop/windows/frontend.go b/v2/internal/frontend/desktop/windows/frontend.go
index fd348394f..dbda34cca 100644
--- a/v2/internal/frontend/desktop/windows/frontend.go
+++ b/v2/internal/frontend/desktop/windows/frontend.go
@@ -19,6 +19,7 @@ import (
 	"github.com/wailsapp/wails/v2/internal/binding"
 	"github.com/wailsapp/wails/v2/internal/frontend"
 	"github.com/wailsapp/wails/v2/internal/frontend/assetserver"
+	"github.com/wailsapp/wails/v2/internal/frontend/desktop/common"
 	"github.com/wailsapp/wails/v2/internal/logger"
 	"github.com/wailsapp/wails/v2/pkg/options"
 )
@@ -335,14 +336,16 @@ func (f *Frontend) processRequest(req *edge.ICoreWebView2WebResourceRequest, arg
 	//Get the request
 	uri, _ := req.GetUri()
 
-	// Translate URI
-	uri = strings.TrimPrefix(uri, "file://wails")
-	if !strings.HasPrefix(uri, "/") {
+	file, match, err := common.TranslateUriToFile(uri, "file", "wails")
+	if err != nil {
+		// TODO Handle errors
+		return
+	} else if !match {
 		return
 	}
 
 	// Load file from asset store
-	content, mimeType, err := f.assets.Load(uri)
+	content, mimeType, err := f.assets.Load(file)
 	if err != nil {
 		return
 	}