wg.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. // Solution to part 9 of the Whispering Gophers code lab.
  2. //
  3. // This program extends part 9.
  4. //
  5. // It connects to the peer specified by -peer.
  6. // It accepts connections from peers and receives messages from them.
  7. // When it sees a peer with an address it hasn't seen before, it makes a
  8. // connection to that peer.
  9. // It adds an ID field containing a random string to each outgoing message.
  10. // When it recevies a message with an ID it hasn't seen before, it broadcasts
  11. // that message to all connected peers.
  12. //
  13. package main
  14. import (
  15. "bufio"
  16. "encoding/json"
  17. "flag"
  18. "fmt"
  19. "log"
  20. "net"
  21. "os"
  22. "sync"
  23. "code.osinet.fr/fgm/whispering_gophers/util"
  24. )
  25. var (
  26. selfAddr = flag.String("self", "", "self host")
  27. peerAddr = flag.String("peer", "", "peer host:port")
  28. self string
  29. )
  30. type Message struct {
  31. ID string
  32. Addr string
  33. Body string
  34. }
  35. func main() {
  36. flag.Parse()
  37. var l net.Listener
  38. var err error
  39. if *selfAddr != "" {
  40. l, err = net.Listen("tcp4", *selfAddr + ":0")
  41. } else {
  42. l, err = util.ListenOnFirstUsableInterface()
  43. }
  44. if err != nil {
  45. log.Fatal(err)
  46. }
  47. self = l.Addr().String()
  48. log.Println("Listening on", self)
  49. go dial(*peerAddr)
  50. go readInput()
  51. for {
  52. c, err := l.Accept()
  53. if err != nil {
  54. log.Fatal(err)
  55. }
  56. go serve(c)
  57. }
  58. }
  59. var peers = &Peers{m: make(map[string]chan<- Message)}
  60. type Peers struct {
  61. m map[string]chan<- Message
  62. mu sync.RWMutex
  63. }
  64. // Add creates and returns a new channel for the given peer address.
  65. // If an address already exists in the registry, it returns nil.
  66. func (p *Peers) Add(addr string) <-chan Message {
  67. p.mu.Lock()
  68. defer p.mu.Unlock()
  69. if _, ok := p.m[addr]; ok {
  70. return nil
  71. }
  72. ch := make(chan Message)
  73. p.m[addr] = ch
  74. return ch
  75. }
  76. // Remove deletes the specified peer from the registry.
  77. func (p *Peers) Remove(addr string) {
  78. p.mu.Lock()
  79. defer p.mu.Unlock()
  80. delete(p.m, addr)
  81. }
  82. // List returns a slice of all active peer channels.
  83. func (p *Peers) List() []chan<- Message {
  84. p.mu.RLock()
  85. defer p.mu.RUnlock()
  86. l := make([]chan<- Message, 0, len(p.m))
  87. for _, ch := range p.m {
  88. l = append(l, ch)
  89. }
  90. return l
  91. }
  92. func broadcast(m Message) {
  93. for _, ch := range peers.List() {
  94. select {
  95. case ch <- m:
  96. default:
  97. // Okay to drop messages sometimes.
  98. }
  99. }
  100. }
  101. func serve(c net.Conn) {
  102. defer c.Close()
  103. d := json.NewDecoder(c)
  104. for {
  105. var m Message
  106. err := d.Decode(&m)
  107. if err != nil {
  108. log.Println(err)
  109. return
  110. }
  111. if Seen(m.ID) {
  112. continue
  113. }
  114. fmt.Printf("%#v\n", m)
  115. broadcast(m)
  116. go dial(m.Addr)
  117. }
  118. }
  119. func readInput() {
  120. s := bufio.NewScanner(os.Stdin)
  121. for s.Scan() {
  122. m := Message{
  123. ID: util.RandomID(),
  124. Addr: self,
  125. Body: s.Text(),
  126. }
  127. Seen(m.ID)
  128. broadcast(m)
  129. }
  130. if err := s.Err(); err != nil {
  131. log.Fatal(err)
  132. }
  133. }
  134. func dial(addr string) {
  135. if addr == self {
  136. return // Don't try to dial self.
  137. }
  138. ch := peers.Add(addr)
  139. if ch == nil {
  140. return // Peer already connected.
  141. }
  142. defer peers.Remove(addr)
  143. c, err := net.Dial("tcp", addr)
  144. if err != nil {
  145. log.Println(addr, err)
  146. return
  147. }
  148. defer c.Close()
  149. e := json.NewEncoder(c)
  150. for m := range ch {
  151. err := e.Encode(m)
  152. if err != nil {
  153. log.Println(addr, err)
  154. return
  155. }
  156. }
  157. }
  158. var seenIDs = struct {
  159. m map[string]bool
  160. sync.Mutex
  161. }{m: make(map[string]bool)}
  162. // Seen returns true if the specified id has been seen before.
  163. // If not, it returns false and marks the given id as "seen".
  164. func Seen(id string) bool {
  165. seenIDs.Lock()
  166. ok := seenIDs.m[id]
  167. seenIDs.m[id] = true
  168. seenIDs.Unlock()
  169. return ok
  170. }