github.com/kvattikuti/drone@v0.2.1-0.20140603034306-d400229a327a/pkg/channel/channel.go (about)

     1  package channel
     2  
     3  import (
     4  	"crypto/rand"
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  	"time"
     9  
    10  	"code.google.com/p/go.net/websocket"
    11  	"github.com/dchest/authcookie"
    12  )
    13  
    14  // secret key used to generate tokens
    15  var secret = make([]byte, 32)
    16  
    17  func init() {
    18  	// generate the secret key by reading
    19  	// from crypto/random
    20  	if _, err := io.ReadFull(rand.Reader, secret); err != nil {
    21  		panic(err)
    22  	}
    23  }
    24  
    25  // Create will generate a token and create a new
    26  // channel over which messages will be sent.
    27  func Create(name string) string {
    28  	mu.Lock()
    29  	defer mu.Unlock()
    30  
    31  	if _, ok := hubs[name]; !ok {
    32  		hub := newHub(false, true)
    33  		hubs[name] = hub
    34  		go hub.run()
    35  	}
    36  	return authcookie.NewSinceNow(name, 24*time.Hour, secret)
    37  }
    38  
    39  // CreateStream will generate a token and create a new
    40  // channel over which messages streams (ie build output)
    41  // are sent.
    42  func CreateStream(name string) string {
    43  	mu.Lock()
    44  	defer mu.Unlock()
    45  
    46  	if _, ok := hubs[name]; !ok {
    47  		hub := newHub(true, false)
    48  		hubs[name] = hub
    49  		go hub.run()
    50  	}
    51  	return authcookie.NewSinceNow(name, 24*time.Hour, secret)
    52  }
    53  
    54  // Token will generate a token, but will not create
    55  // a new channel.
    56  func Token(name string) string {
    57  	return authcookie.NewSinceNow(name, 24*time.Hour, secret)
    58  }
    59  
    60  // Send sends a message on the named channel.
    61  func Send(name string, message string) error {
    62  	return SendBytes(name, []byte(message))
    63  }
    64  
    65  // SendJSON sends a JSON-encoded value on
    66  // the named channel.
    67  func SendJSON(name string, value interface{}) error {
    68  	m, err := json.Marshal(value)
    69  	if err != nil {
    70  		return err
    71  	}
    72  
    73  	return SendBytes(name, m)
    74  }
    75  
    76  // SendBytes send a message in byte format on
    77  // the named channel.
    78  func SendBytes(name string, value []byte) error {
    79  	// get the hub for the specified channel name
    80  	mu.RLock()
    81  	hub, ok := hubs[name]
    82  	mu.RUnlock()
    83  
    84  	if !ok {
    85  		return fmt.Errorf("channel does not exist")
    86  	}
    87  
    88  	go hub.Write(value)
    89  	return nil
    90  }
    91  
    92  func Read(ws *websocket.Conn) {
    93  
    94  	// get the name from the request
    95  	hash := ws.Request().FormValue("token")
    96  
    97  	// get the hash of the token
    98  	name := authcookie.Login(hash, secret)
    99  
   100  	// get the hub for the specified channel name
   101  	mu.RLock()
   102  	hub, ok := hubs[name]
   103  	mu.RUnlock()
   104  
   105  	// if hub not found, exit
   106  	if !ok {
   107  		ws.Close()
   108  		return
   109  	}
   110  
   111  	// internal representation of a connection
   112  	// maximum queue of 100000 messages
   113  	conn := &connection{
   114  		send: make(chan string, 100000),
   115  		ws:   ws,
   116  	}
   117  
   118  	// register the connection with the hub
   119  	hub.register <- conn
   120  
   121  	defer func() {
   122  		go func() {
   123  			hub.unregister <- conn
   124  		}()
   125  		closed := <-hub.closed
   126  
   127  		// this will remove the hub when the connection is
   128  		// closed if the
   129  		if hub.autoClose && closed {
   130  			mu.Lock()
   131  			delete(hubs, name)
   132  			mu.Unlock()
   133  		}
   134  	}()
   135  
   136  	go conn.writer()
   137  	conn.reader()
   138  }
   139  
   140  func Close(name string) {
   141  	// get the hub for the specified channel name
   142  	mu.RLock()
   143  	hub, ok := hubs[name]
   144  	mu.RUnlock()
   145  
   146  	if !ok {
   147  		return
   148  	}
   149  
   150  	// close hub connections
   151  	hub.Close()
   152  
   153  	// remove the hub
   154  	mu.Lock()
   155  	delete(hubs, name)
   156  	mu.Unlock()
   157  }