11package kafka
22
33import (
4+ "context"
5+ "crypto/tls"
46 "time"
57
68 "github.com/moov-io/achgateway/internal/service"
79 "github.com/moov-io/base/log"
810
911 "github.com/Shopify/sarama"
12+ "github.com/aws/aws-msk-iam-sasl-signer-go/signer"
1013 "gocloud.dev/pubsub"
1114 "gocloud.dev/pubsub/kafkapubsub"
1215)
@@ -15,30 +18,66 @@ var (
1518 minKafkaVersion = sarama .V2_6_0_0
1619)
1720
21+ type MSKAccessTokenProvider struct {
22+ Region string
23+ Profile string
24+ RoleARN string
25+ SessionName string
26+ }
27+
28+ func (m * MSKAccessTokenProvider ) Token () (* sarama.AccessToken , error ) {
29+ var token string
30+ var err error
31+
32+ // Choose the correct AWS authentication method
33+ switch {
34+ case m .Profile != "" :
35+ token , _ , err = signer .GenerateAuthTokenFromProfile (context .TODO (), m .Region , m .Profile )
36+ case m .RoleARN != "" :
37+ token , _ , err = signer .GenerateAuthTokenFromRole (context .TODO (), m .Region , m .RoleARN , m .SessionName )
38+ default :
39+ token , _ , err = signer .GenerateAuthToken (context .TODO (), m .Region )
40+ }
41+
42+ if err != nil {
43+ return nil , err
44+ }
45+ return & sarama.AccessToken {Token : token }, nil
46+ }
47+
1848func OpenTopic (logger log.Logger , cfg * service.KafkaConfig ) (* pubsub.Topic , error ) {
1949 config := kafkapubsub .MinimalConfig ()
2050 config .Version = minKafkaVersion
2151 config .Net .TLS .Enable = cfg .TLS
2252
2353 config .Net .SASL .Enable = cfg .Key != ""
24- config .Net .SASL .Mechanism = sarama .SASLMechanism (cfg .SASLMechanism )
25-
2654 // Default to PLAIN if no SASL mechanism is specified
2755 switch cfg .SASLMechanism {
2856 case "SCRAM-SHA-512" :
2957 config .Net .SASL .SCRAMClientGeneratorFunc = func () sarama.SCRAMClient {
3058 return & XDGSCRAMClient {HashGeneratorFcn : SHA512 }
3159 }
32- config .Net .SASL .Mechanism = sarama .SASLMechanism ( cfg . SASLMechanism )
60+ config .Net .SASL .Mechanism = sarama .SASLTypeSCRAMSHA512
3361
3462 case "SCRAM-SHA-256" :
3563 config .Net .SASL .SCRAMClientGeneratorFunc = func () sarama.SCRAMClient {
3664 return & XDGSCRAMClient {HashGeneratorFcn : SHA256 }
3765 }
38- config .Net .SASL .Mechanism = sarama .SASLMechanism (cfg .SASLMechanism )
66+ config .Net .SASL .Mechanism = sarama .SASLTypeSCRAMSHA256
67+
68+ case "AWS_MSK_IAM" :
69+ config .Net .SASL .Mechanism = sarama .SASLTypeOAuth
70+ config .Net .SASL .TokenProvider = & MSKAccessTokenProvider {
71+ Region : cfg .AWSRegion ,
72+ Profile : cfg .AWSProfile ,
73+ RoleARN : cfg .AWSRoleARN ,
74+ SessionName : cfg .AWSSessionName ,
75+ }
76+ config .Net .TLS .Enable = true
77+ config .Net .TLS .Config = & tls.Config {}
3978
4079 default :
41- config .Net .SASL .Mechanism = sarama .SASLMechanism ( "PLAIN" )
80+ config .Net .SASL .Mechanism = sarama .SASLTypePlaintext
4281 }
4382
4483 config .Net .SASL .User = cfg .Key
@@ -52,8 +91,10 @@ func OpenTopic(logger log.Logger, cfg *service.KafkaConfig) (*pubsub.Topic, erro
5291 Set ("tls" , log .Bool (cfg .TLS )).
5392 Set ("group" , log .String (cfg .Group )).
5493 Set ("sasl.enable" , log .Bool (config .Net .SASL .Enable )).
94+ Set ("sasl.mechanism" , log .String (string (config .Net .SASL .Mechanism ))).
5595 Set ("sasl.user" , log .String (cfg .Key )).
5696 Set ("topic" , log .String (cfg .Topic )).
97+ Set ("aws.region" , log .String (cfg .AWSRegion )).
5798 Log ("opening kafka topic" )
5899
59100 return kafkapubsub .OpenTopic (cfg .Brokers , config , cfg .Topic , nil )
@@ -71,16 +112,27 @@ func OpenSubscription(logger log.Logger, cfg *service.KafkaConfig) (*pubsub.Subs
71112 config .Net .SASL .SCRAMClientGeneratorFunc = func () sarama.SCRAMClient {
72113 return & XDGSCRAMClient {HashGeneratorFcn : SHA512 }
73114 }
74- config .Net .SASL .Mechanism = sarama .SASLMechanism ( cfg . SASLMechanism )
115+ config .Net .SASL .Mechanism = sarama .SASLTypeSCRAMSHA512
75116
76117 case "SCRAM-SHA-256" :
77118 config .Net .SASL .SCRAMClientGeneratorFunc = func () sarama.SCRAMClient {
78119 return & XDGSCRAMClient {HashGeneratorFcn : SHA256 }
79120 }
80- config .Net .SASL .Mechanism = sarama .SASLMechanism (cfg .SASLMechanism )
121+ config .Net .SASL .Mechanism = sarama .SASLTypeSCRAMSHA256
122+
123+ case "AWS_MSK_IAM" :
124+ config .Net .SASL .Mechanism = sarama .SASLTypeOAuth
125+ config .Net .SASL .TokenProvider = & MSKAccessTokenProvider {
126+ Region : cfg .AWSRegion ,
127+ Profile : cfg .AWSProfile ,
128+ RoleARN : cfg .AWSRoleARN ,
129+ SessionName : cfg .AWSSessionName ,
130+ }
131+ config .Net .TLS .Enable = true
132+ config .Net .TLS .Config = & tls.Config {}
81133
82134 default :
83- config .Net .SASL .Mechanism = sarama .SASLMechanism ( "PLAIN" )
135+ config .Net .SASL .Mechanism = sarama .SASLTypePlaintext
84136 }
85137
86138 config .Net .SASL .User = cfg .Key
@@ -98,8 +150,10 @@ func OpenSubscription(logger log.Logger, cfg *service.KafkaConfig) (*pubsub.Subs
98150 Set ("tls" , log .Bool (cfg .TLS )).
99151 Set ("group" , log .String (cfg .Group )).
100152 Set ("sasl.enable" , log .Bool (config .Net .SASL .Enable )).
153+ Set ("sasl.mechanism" , log .String (string (config .Net .SASL .Mechanism ))).
101154 Set ("sasl.user" , log .String (cfg .Key )).
102155 Set ("topic" , log .String (cfg .Topic )).
156+ Set ("aws.region" , log .String (cfg .AWSRegion )).
103157 Log ("setting up kafka subscription" )
104158
105159 return kafkapubsub .OpenSubscription (cfg .Brokers , config , cfg .Group , []string {cfg .Topic }, & kafkapubsub.SubscriptionOptions {
0 commit comments