github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/modules/sock/sock.go (about)

     1  package sock
     2  
     3  import (
     4  	//"errors"
     5  	"fmt"
     6  	"github.com/insionng/makross"
     7  	"log"
     8  	"net"
     9  	"net/http"
    10  	"strings"
    11  	"sync"
    12  	"time"
    13  	"github.com/insionng/yougam/helper"
    14  	"github.com/insionng/yougam/libraries/gorilla/websocket"
    15  	"github.com/insionng/yougam/models"
    16  
    17  	"github.com/insionng/yougam/modules/setting"
    18  	"github.com/insionng/yougam/routers"
    19  )
    20  
    21  type Client struct {
    22  	websocket *websocket.Conn
    23  	clientIP  net.Addr
    24  	userid    int64
    25  }
    26  
    27  type Box struct {
    28  	sync.RWMutex
    29  	key string
    30  	box []*Client
    31  }
    32  
    33  type Boxes struct {
    34  	sync.RWMutex
    35  	boxes []*Box
    36  }
    37  
    38  func (b *Box) appendClient(client *Client) {
    39  	b.Lock()
    40  	b.box = append(b.box, client)
    41  	for _, c := range b.box {
    42  		if c != client {
    43  			log.Println(client.clientIP, "进入聊天~")
    44  		}
    45  	}
    46  	b.Unlock()
    47  }
    48  
    49  func (b *Box) removeClient(client *Client) {
    50  	b.Lock()
    51  	defer b.Unlock()
    52  
    53  	for index, c := range b.box {
    54  		if c == client {
    55  			b.box = append(b.box[:index], b.box[(index+1):]...)
    56  			if err := c.websocket.WriteMessage(1, []byte(`<div class="alert alert-warning alert-dismissible" role="alert"><button type="button" class="close" data-dismiss="alert"><span aria-hidden="true">&times;</span><span class="sr-only">Close</span></button><h4 class="alert-heading">警告</h4>`+`你的账号在别处连接,当前通信已经失效,如非本人操作请立即修改账号密码..`+`</div>`)); err != nil {
    57  				return
    58  			}
    59  		} else {
    60  			if err := client.websocket.WriteMessage(1, []byte(`<div class="alert alert-warning alert-dismissible" role="alert"><button type="button" class="close" data-dismiss="alert"><span aria-hidden="true">&times;</span><span class="sr-only">Close</span></button><h4 class="alert-heading">警告</h4>`+`你的账号已经重新连接,原连接已经失效,系统禁止同一账号在多客户端同时连接..`+`</div>`)); err != nil {
    61  				return
    62  			}
    63  		}
    64  	}
    65  }
    66  
    67  func (b *Box) removeUser(client *Client) {
    68  	b.Lock()
    69  	defer b.Unlock()
    70  
    71  	for index, c := range b.box {
    72  		if c.userid == client.userid {
    73  			b.box = append(b.box[:index], b.box[(index+1):]...)
    74  			if err := c.websocket.WriteMessage(1, []byte(`<div class="alert alert-warning alert-dismissible" role="alert"><button type="button" class="close" data-dismiss="alert"><span aria-hidden="true">&times;</span><span class="sr-only">Close</span></button><h4 class="alert-heading">警告</h4>`+`你的账号在别处连接,当前通信已经失效,如非本人操作请立即修改账号密码..`+`</div>`)); err != nil {
    75  				return
    76  			}
    77  		} else {
    78  			if err := client.websocket.WriteMessage(1, []byte(`<div class="alert alert-warning alert-dismissible" role="alert"><button type="button" class="close" data-dismiss="alert"><span aria-hidden="true">&times;</span><span class="sr-only">Close</span></button><h4 class="alert-heading">警告</h4>`+`你的账号已经重新连接,原连接已经失效,系统禁止同一账号在多客户端同时连接..`+`</div>`)); err != nil {
    79  				return
    80  			}
    81  		}
    82  	}
    83  }
    84  
    85  func (b *Box) broadcastMessage(messageType int, message []byte) {
    86  	b.Lock()
    87  	defer b.Unlock()
    88  
    89  	for _, client := range b.box {
    90  		if err := client.websocket.WriteMessage(messageType, message); err != nil {
    91  			return
    92  		}
    93  	}
    94  }
    95  
    96  func (b *Boxes) getBox(key string) *Box {
    97  	b.Lock()
    98  	defer b.Unlock()
    99  
   100  	for _, Box := range b.boxes {
   101  		if Box.key == key {
   102  			return Box
   103  		}
   104  	}
   105  
   106  	box := &Box{sync.RWMutex{}, key, make([]*Client, 0)}
   107  	b.boxes = append(b.boxes, box)
   108  	return box
   109  }
   110  
   111  func newBoxes() *Boxes {
   112  	return &Boxes{sync.RWMutex{}, make([]*Box, 0)}
   113  }
   114  
   115  func Sock(m *makross.Macross) {
   116  
   117  	if setting.FlashPolicyService {
   118  		go FlashPolicyService()
   119  	}
   120  
   121  	var boxes *Boxes
   122  	boxes = newBoxes()
   123  
   124  	m.Get("/sock/connect/:name([\\x{4e00}-\\x{9fa5}A-Z0-9a-z_-]+)/", func(self *self.Context) {
   125  		if name := self.ParamsEscape("name"); self.IsSigned && (len(name) > 0) && websocket.IsWebSocketUpgrade(self.Req.Request) {
   126  			recipient, err := models.GetUserByUsername(name)
   127  			if (err != nil) || (recipient == nil) {
   128  				log.Println(err)
   129  				return // err
   130  			}
   131  
   132  			ws, err := websocket.Upgrade(self.Resp, self.Req.Request, nil, 1024, 1024)
   133  			if _, okay := err.(websocket.HandshakeError); okay {
   134  				log.Println("Not a websocket handshake")
   135  				http.Error(self.Resp, "Not a websocket handshake", 400)
   136  				return // err
   137  			} else if err != nil {
   138  				log.Println(err)
   139  				return // err
   140  			}
   141  
   142  			rAddr := ws.RemoteAddr()
   143  			sockCli := &Client{ws, rAddr, self.User.Id}
   144  
   145  			box := boxes.getBox(helper.OrderKey(self.User.Id, recipient.Id))
   146  
   147  			if len(box.box) < 2 {
   148  				log.Println("L<2")
   149  				box.appendClient(sockCli)
   150  			} else {
   151  				log.Println("L>=2")
   152  				uid := fmt.Sprintf("%v", self.User.Id)
   153  				//如果是原来的用户则更新连接地址,即禁止用户同时登录2个以上客户端
   154  				if s := strings.Split(box.key, ":"); (s[0] == uid) || (s[1] == uid) {
   155  					box.removeUser(sockCli)
   156  					box.appendClient(sockCli)
   157  				} else {
   158  					return // nil
   159  				}
   160  
   161  			}
   162  
   163  			for {
   164  
   165  				messageType, message, err := ws.ReadMessage()
   166  				if err != nil {
   167  					box.removeClient(sockCli)
   168  					//log.Println("bye")
   169  					log.Println(err)
   170  					return // err
   171  				}
   172  
   173  				now := time.Now()
   174  				policy := helper.ObjPolicy()
   175  				s := policy.SanitizeBytes(message)
   176  				if len(self.User.AvatarMedium) <= 0 {
   177  					self.User.AvatarMedium = "/img/d48.png"
   178  				}
   179  
   180  				body := fmt.Sprintf(`
   181                      <article class="panel panel-warning box cell first message">
   182                          <div class="panel-body">%s</div>
   183                          <div id="uid%v" class="panel-footer box nobg cell last">
   184  				            <a href="/user/%s/" class="thumb avatar">
   185  				            	<p><img src="%s"></p><p>%s</p>
   186  				            </a>
   187  	                        <div style="clear:both;"></div>
   188  	                        <small class="pull-right">%s</small>
   189                          </div>
   190                          <div style="clear:both;"></div>
   191                      </article>`, s,
   192  					self.User.Id, self.User.Username,
   193  					self.User.AvatarMedium, self.User.Username,
   194  					now.Format("2006.01.02 03:04"))
   195  				box.broadcastMessage(messageType, []byte(body))
   196  
   197  			}
   198  		} else {
   199  			return // errors.New("Is not WebSocket Upgrade")
   200  		}
   201  	})
   202  
   203  	m.Get("/contact/", routers.GetContactHandler)
   204  	m.Get("/contact/search/", routers.GetContactHandler)
   205  	m.Get("/connect/<name:([\\x{4e00}-\\x{9fa5}A-Z0-9a-z_-]+)>/", routers.GetContactHandler)
   206  
   207  }