go: only allow origins from same network (localhost)

master
ALI Hamza 2019-10-13 10:14:33 +07:00
parent 32ae99c225
commit 0cb7f55d63
No known key found for this signature in database
GPG Key ID: 7C608266BF384ADC
2 changed files with 7 additions and 2 deletions

@ -1,13 +1,13 @@
package main package main
import ( import (
"TerraOceanBot/server"
"io/ioutil" "io/ioutil"
"os" "os"
"os/signal" "os/signal"
"syscall" "syscall"
"TerraOceanBot/discord" "TerraOceanBot/discord"
"TerraOceanBot/server"
) )
func main() { func main() {

@ -5,6 +5,7 @@ import (
"log" "log"
"net/http" "net/http"
"strconv" "strconv"
"strings"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )
@ -12,12 +13,16 @@ import (
var upgrader = websocket.Upgrader{ var upgrader = websocket.Upgrader{
ReadBufferSize: 1024, ReadBufferSize: 1024,
WriteBufferSize: 1024, WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool {
return strings.Contains(r.Host, "localhost")
},
} }
func newConnection(w http.ResponseWriter, r *http.Request) { func newConnection(w http.ResponseWriter, r *http.Request) {
ws, err := upgrader.Upgrade(w, r, nil) ws, err := upgrader.Upgrade(w, r, nil)
if err != 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() defer ws.Close()