func (cli *Client) handshake(c diam.Conn) (diam.Conn, error) { ip, _, err := net.SplitHostPort(c.LocalAddr().String()) if err != nil { return nil, err } m := cli.makeCER(net.ParseIP(ip)) // Ignore CER, but not DWR. cli.Handler.mux.HandleFunc("CER", func(c diam.Conn, m *diam.Message) {}) // Handle CEA and DWA. errc := make(chan error) dwac := make(chan struct{}) cli.Handler.mux.Handle("CEA", handleCEA(cli.Handler, errc)) cli.Handler.mux.Handle("DWA", handshakeOK(handleDWA(cli.Handler, dwac))) for i := 0; i < (int(cli.MaxRetransmits) + 1); i++ { _, err := m.WriteTo(c) if err != nil { return nil, err } select { case err := <-errc: // Wait for CEA. if err != nil { close(errc) return nil, err } if cli.EnableWatchdog { go cli.watchdog(c, dwac) } return c, nil case <-time.After(cli.RetransmitInterval): } } c.Close() return nil, ErrHandshakeTimeout }
func (cli *Client) dwr(c diam.Conn, osid uint32, dwac chan struct{}) { m := cli.makeDWR(osid) for i := 0; i < (int(cli.MaxRetransmits) + 1); i++ { _, err := m.WriteTo(c) if err != nil { return } select { case <-dwac: return case <-time.After(cli.RetransmitInterval): } } // Watchdog failed, disconnect. c.Close() }
// Pump messages from one side to the other. func Pump(src, dst diam.Conn, srcChan, dstChan chan *diam.Message) { for { select { case m := <-srcChan: if m == nil { src.Close() return } log.Printf( "Message from %s to %s\n%s", src.RemoteAddr().String(), dst.RemoteAddr().String(), m, ) if _, err := m.WriteTo(src); err != nil { src.Close() // triggers the case below } case <-src.(diam.CloseNotifier).CloseNotify(): liveMu.Lock() defer liveMu.Unlock() if _, ok := liveBridge[src.RemoteAddr().String()]; ok { delete(liveBridge, src.RemoteAddr().String()) log.Printf( "Destroying bridge from %s to %s", src.RemoteAddr().String(), dst.RemoteAddr().String(), ) } else { delete(liveBridge, dst.RemoteAddr().String()) log.Printf( "Destroying bridge from %s to %s", dst.RemoteAddr().String(), src.RemoteAddr().String(), ) } src.Close() dstChan <- nil return } } }