diff --git a/src/github.com/matrix-org/dendrite/cmd/dendrite-room-server/main.go b/src/github.com/matrix-org/dendrite/cmd/dendrite-room-server/main.go index f607d1ec2..06773972a 100644 --- a/src/github.com/matrix-org/dendrite/cmd/dendrite-room-server/main.go +++ b/src/github.com/matrix-org/dendrite/cmd/dendrite-room-server/main.go @@ -20,7 +20,6 @@ import ( _ "net/http/pprof" "os" - log "github.com/sirupsen/logrus" "github.com/matrix-org/dendrite/common" "github.com/matrix-org/dendrite/common/config" "github.com/matrix-org/dendrite/roomserver/alias" @@ -28,6 +27,7 @@ import ( "github.com/matrix-org/dendrite/roomserver/query" "github.com/matrix-org/dendrite/roomserver/storage" "github.com/prometheus/client_golang/prometheus" + log "github.com/sirupsen/logrus" sarama "gopkg.in/Shopify/sarama.v1" ) diff --git a/src/github.com/matrix-org/dendrite/roomserver/query/query.go b/src/github.com/matrix-org/dendrite/roomserver/query/query.go index c897a9b4e..bddc9cf86 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/query/query.go +++ b/src/github.com/matrix-org/dendrite/roomserver/query/query.go @@ -16,8 +16,10 @@ package query import ( "context" + "database/sql" "encoding/json" "net/http" + "time" "github.com/matrix-org/dendrite/common" "github.com/matrix-org/dendrite/roomserver/api" @@ -74,6 +76,10 @@ type RoomserverQueryAPIDatabase interface { EventStateKeys( context.Context, []types.EventStateKeyNID, ) (map[types.EventStateKeyNID]string, error) + // Lookup if the roomID has been reserved, and when that reservation was made. + RoomIDReserved(ctx context.Context, roomID string) (time.Time, error) + // Reserve the room ID. + ReserveRoomID(ctx context.Context, roomID string) error } // RoomserverQueryAPI is an implementation of api.RoomserverQueryAPI @@ -425,14 +431,34 @@ func (r *RoomserverQueryAPI) QueryReserveRoomID( response *api.QueryReserveRoomIDResponse, ) error { nid, err := r.DB.RoomNID(ctx, request.RoomID) + if err != nil && err != sql.ErrNoRows { + return err + } + if nid != 0 { response.Success = false return nil } + // Check if Room ID has already been reserved. + reservedSince, err := r.DB.RoomIDReserved(ctx, request.RoomID) + if reservedSince != (time.Time{}) { + response.Success = false + return nil + } + + if err != nil && err != sql.ErrNoRows { + return err + } + + err = r.DB.ReserveRoomID(ctx, request.RoomID) + if err != nil { + return err + } + response.Success = true - return err + return nil } // SetupHTTP adds the RoomserverQueryAPI handlers to the http.ServeMux. diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go b/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go index a24dbb1d3..1a8f497f5 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go +++ b/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go @@ -30,6 +30,7 @@ type statements struct { roomAliasesStatements inviteStatements membershipStatements + reservedRoomStatements } func (s *statements) prepare(db *sql.DB) error { @@ -47,6 +48,7 @@ func (s *statements) prepare(db *sql.DB) error { s.roomAliasesStatements.prepare, s.inviteStatements.prepare, s.membershipStatements.prepare, + s.reservedRoomStatements.prepare, } { if err = prepare(db); err != nil { return err diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go index ad4fed659..ee812710a 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go +++ b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go @@ -17,6 +17,7 @@ package storage import ( "context" "database/sql" + "time" // Import the postgres database driver. _ "github.com/lib/pq" @@ -651,6 +652,17 @@ func (d *Database) GetMembershipEventNIDsForRoom( return d.statements.selectMembershipsFromRoom(ctx, roomNID) } +// RoomIDReserved implements query.RoomserverQueryAPIDB +func (d *Database) RoomIDReserved(ctx context.Context, roomID string) (time.Time, error) { + return d.statements.selectReservedRoom(ctx, roomID) +} + +// ReserveRoomID implements query.RoomserverQueryAPIDB +// It saved the given room ID as reserved. +func (d *Database) ReserveRoomID(ctx context.Context, roomID string) error { + return d.statements.insertReservedRoom(ctx, roomID, time.Now()) +} + type transaction struct { ctx context.Context txn *sql.Tx