Add Relay Info Forwarding for Snowflake

2 files changed, 26 insertions, 1 deletions

diff --git a/broker/broker.go b/broker/broker.go
index 10129d7..692cea4 100644
--- a/broker/broker.go
+++ b/broker/broker.go
@@ -36,6 +36,13 @@ type BrokerContext struct {
 	snowflakeLock sync.Mutex
 	proxyPolls    chan *ProxyPoll
 	metrics       *Metrics
+
+	bridgeList          BridgeListHolderFileBased
+	allowedRelayPattern string
+}
+
+func (ctx *BrokerContext) GetBridgeInfo(fingerprint [20]byte) (BridgeInfo, error) {
+	return ctx.bridgeList.GetBridgeInfo(fingerprint)
 }
 
 func NewBrokerContext(metricsLogger *log.Logger) *BrokerContext {
@@ -139,6 +146,14 @@ func (ctx *BrokerContext) AddSnowflake(id string, proxyType string, natType stri
 	return snowflake
 }
 
+func (ctx *BrokerContext) InstallBridgeListProfile(reader io.Reader, relayPattern string) error {
+	if err := ctx.bridgeList.LoadBridgeInfo(reader); err != nil {
+		return err
+	}
+	ctx.allowedRelayPattern = relayPattern
+	return nil
+}
+
 // Client offer contains an SDP, bridge fingerprint and the NAT type of the client
 type ClientOffer struct {
 	natType     string
diff --git a/broker/ipc.go b/broker/ipc.go
index e11a33c..e559c2a 100644
--- a/broker/ipc.go
+++ b/broker/ipc.go
@@ -102,7 +102,13 @@ func (i *IPC) ProxyPolls(arg messages.Arg, response *[]byte) error {
 	}
 
 	i.ctx.metrics.promMetrics.ProxyPollTotal.With(prometheus.Labels{"nat": natType, "status": "matched"}).Inc()
-	b, err = messages.EncodePollResponse(string(offer.sdp), true, offer.natType)
+	var relayURL string
+	if info, err := i.ctx.bridgeList.GetBridgeInfo(offer.fingerprint); err != nil {
+		return err
+	} else {
+		relayURL = info.WebSocketAddress
+	}
+	b, err = messages.EncodePollResponseWithRelayURL(string(offer.sdp), true, offer.natType, relayURL)
 	if err != nil {
 		return messages.ErrInternal
 	}
@@ -141,6 +147,10 @@ func (i *IPC) ClientOffers(arg messages.Arg, response *[]byte) error {
 	}
 	copy(offer.fingerprint[:], fingerprint)
 
+	if _, err := i.ctx.GetBridgeInfo(offer.fingerprint); err != nil {
+		return err
+	}
+
 	// Only hand out known restricted snowflakes to unrestricted clients
 	var snowflakeHeap *SnowflakeHeap
 	if offer.natType == NATUnrestricted {