diff --git a/GoBot/main.go b/GoBot/main.go
index 311c492..59c0f4c 100644
--- a/GoBot/main.go
+++ b/GoBot/main.go
@@ -1,13 +1,13 @@
 package main
 
 import (
-	"TerraOceanBot/server"
 	"io/ioutil"
 	"os"
 	"os/signal"
 	"syscall"
 
 	"TerraOceanBot/discord"
+	"TerraOceanBot/server"
 )
 
 func main() {
diff --git a/GoBot/server/server.go b/GoBot/server/server.go
index 6f521ac..7d207cf 100644
--- a/GoBot/server/server.go
+++ b/GoBot/server/server.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"net/http"
 	"strconv"
+	"strings"
 
 	"github.com/gorilla/websocket"
 )
@@ -12,12 +13,16 @@ import (
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return strings.Contains(r.Host, "localhost")
+	},
 }
 
 func newConnection(w http.ResponseWriter, r *http.Request) {
 	ws, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
-		fmt.Fprint(w, "You must use the websocket protocol to connect to this endpoint.", err)
+		fmt.Fprintf(w, "An error has occured: %s.\n", err)
+		return
 	}
 	defer ws.Close()