gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/agent/input/discord/conn.go (about)

     1  package discord
     2  
     3  import (
     4  	"errors"
     5  	"strings"
     6  	"sync"
     7  
     8  	"github.com/bwmarrin/discordgo"
     9  	"gitee.com/liuxuezhan/go-micro-v1.18.0/agent/input"
    10  	"gitee.com/liuxuezhan/go-micro-v1.18.0/util/log"
    11  )
    12  
    13  type discordConn struct {
    14  	master *discordInput
    15  	exit   chan struct{}
    16  	recv   chan *discordgo.Message
    17  
    18  	sync.Mutex
    19  }
    20  
    21  func newConn(master *discordInput) *discordConn {
    22  	conn := &discordConn{
    23  		master: master,
    24  		exit:   make(chan struct{}),
    25  		recv:   make(chan *discordgo.Message),
    26  	}
    27  
    28  	conn.master.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
    29  		if m.Author.ID == master.botID {
    30  			return
    31  		}
    32  
    33  		whitelisted := false
    34  		for _, ID := range conn.master.whitelist {
    35  			if m.Author.ID == ID {
    36  				whitelisted = true
    37  				break
    38  			}
    39  		}
    40  
    41  		if len(master.whitelist) > 0 && !whitelisted {
    42  			return
    43  		}
    44  
    45  		var valid bool
    46  		m.Message.Content, valid = conn.master.prefixfn(m.Message.Content)
    47  		if !valid {
    48  			return
    49  		}
    50  
    51  		conn.recv <- m.Message
    52  	})
    53  
    54  	return conn
    55  }
    56  
    57  func (dc *discordConn) Recv(event *input.Event) error {
    58  	for {
    59  		select {
    60  		case <-dc.exit:
    61  			return errors.New("connection closed")
    62  		case msg := <-dc.recv:
    63  
    64  			event.From = msg.ChannelID + ":" + msg.Author.ID
    65  			event.To = dc.master.botID
    66  			event.Type = input.TextEvent
    67  			event.Data = []byte(msg.Content)
    68  			return nil
    69  		}
    70  	}
    71  }
    72  
    73  func (dc *discordConn) Send(e *input.Event) error {
    74  	fields := strings.Split(e.To, ":")
    75  	_, err := dc.master.session.ChannelMessageSend(fields[0], string(e.Data))
    76  	if err != nil {
    77  		log.Log("[bot][loop][send]", err)
    78  	}
    79  	return nil
    80  }
    81  
    82  func (dc *discordConn) Close() error {
    83  	if err := dc.master.session.Close(); err != nil {
    84  		return err
    85  	}
    86  
    87  	select {
    88  	case <-dc.exit:
    89  		return nil
    90  	default:
    91  		close(dc.exit)
    92  	}
    93  	return nil
    94  }