Repository: nats-io/nats-server
Pull request #3325 (2022-08-03 01:39:59)

Allow direct get by subject to be all subject based.
This avoids marshaling or un-marshaling but also allows subject based permissioning.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index e5849b51..1843bc1d 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -124,6 +124,11 @@ const (
 	JSDirectMsgGet  = "$JS.API.DIRECT.GET.*"
 	JSDirectMsgGetT = "$JS.API.DIRECT.GET.%s"
 
+	// This is a direct version of get last by subject, which will be the dominant pattern for KV access once 2.9 is released.
+	// The stream and the key will be part of the subject to allow for no-marshal payloads and subject based security permissions.
+	JSDirectGetLastBySubject  = "$JS.API.DIRECT.GET.*.>"
+	JSDirectGetLastBySubjectT = "$JS.API.DIRECT.GET.%s.%s"
+
 	// jsDirectGetPre
 	jsDirectGetPre = "$JS.API.DIRECT.GET"
 
diff --git a/server/stream.go b/server/stream.go
index 2306782d..62f6cdb9 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -237,6 +237,7 @@ type stream struct {
 
 	// Direct get subscription.
 	directSub *subscription
+	lastBySub *subscription
 }
 
 type sourceInfo struct {
@@ -245,6 +246,7 @@ type sourceInfo struct {
 	cname string
 	sub   *subscription
 	dsub  *subscription
+	lbsub *subscription
 	msgs  *ipQueue // of *inMsg
 	sseq  uint64
 	dseq  uint64
@@ -1433,14 +1435,20 @@ func (mset *stream) updateWithAdvisory(config *StreamConfig, sendAdvisory bool)
 	}
 
 	// Check for mirror. If set but we are not a mirror just ignore for now.
-	if mset.cfg.MirrorDirect && mset.cfg.Mirror != nil && mset.mirror.dsub == nil {
+	if mset.cfg.MirrorDirect && mset.cfg.Mirror != nil {
 		if err := mset.subscribeToMirrorDirect(); err != nil {
 			// Disable since we had problems above.
 			mset.cfg.MirrorDirect = false
 		}
-	} else if !mset.cfg.MirrorDirect && mset.mirror != nil && mset.mirror.dsub != nil {
-		mset.unsubscribe(mset.mirror.dsub)
-		mset.mirror.dsub = nil
+	} else if !mset.cfg.MirrorDirect && mset.mirror != nil {
+		if mset.mirror.dsub != nil {
+			mset.unsubscribe(mset.mirror.dsub)
+			mset.mirror.dsub = nil
+		}
+		if mset.mirror.lbsub != nil {
+			mset.unsubscribe(mset.mirror.lbsub)
+			mset.mirror.lbsub = nil
+		}
 	}
 	mset.mu.Unlock()
 
@@ -2910,31 +2918,56 @@ func (mset *stream) subscribeToStream() error {
 
 // Lock should be held.
 func (mset *stream) subscribeToDirect() error {
-	if mset.directSub != nil {
-		return nil
-	}
-	dsubj := fmt.Sprintf(JSDirectMsgGetT, mset.cfg.Name)
 	// We will make this listen on a queue group by default, which can allow mirrors to participate on opt-in basis.
-	if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetRequest); err != nil {
-		return err
-	} else {
-		mset.directSub = sub
+	if mset.directSub == nil {
+		dsubj := fmt.Sprintf(JSDirectMsgGetT, mset.cfg.Name)
+		if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetRequest); err == nil {
+			mset.directSub = sub
+		} else {
+			return err
+		}
+	}
+	// Now the one that will have subject appended past stream name.
+	if mset.lastBySub == nil {
+		dsubj := fmt.Sprintf(JSDirectGetLastBySubjectT, mset.cfg.Name, fwcs)
+		// We will make this listen on a queue group by default, which can allow mirrors to participate on opt-in basis.
+		if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetLastBySubjectRequest); err == nil {
+			mset.lastBySub = sub
+		} else {
+			return err
+		}
 	}
+
 	return nil
 }
 
 // Lock should be held.
 func (mset *stream) subscribeToMirrorDirect() error {
-	if mset.mirror == nil || mset.mirror.dsub != nil {
+	if mset.mirror == nil {
 		return nil
 	}
-	dsubj := fmt.Sprintf(JSDirectMsgGetT, mset.mirror.name)
+
 	// We will make this listen on a queue group by default, which can allow mirrors to participate on opt-in basis.
-	if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetRequest); err != nil {
-		return err
-	} else {
-		mset.mirror.dsub = sub
+	if mset.mirror.dsub == nil {
+		dsubj := fmt.Sprintf(JSDirectMsgGetT, mset.mirror.name)
+		// We will make this listen on a queue group by default, which can allow mirrors to participate on opt-in basis.
+		if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetRequest); err == nil {
+			mset.mirror.dsub = sub
+		} else {
+			return err
+		}
 	}
+	// Now the one that will have subject appended past stream name.
+	if mset.mirror.lbsub == nil {
+		dsubj := fmt.Sprintf(JSDirectGetLastBySubjectT, mset.mirror.name, fwcs)
+		// We will make this listen on a queue group by default, which can allow mirrors to participate on opt-in basis.
+		if sub, err := mset.queueSubscribeInternal(dsubj, dgetGroup, mset.processDirectGetLastBySubjectRequest); err == nil {
+			mset.mirror.lbsub = sub
+		} else {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -2969,10 +3002,16 @@ func (mset *stream) unsubscribeToStream(stopping bool) error {
 		mset.stopSourceConsumers()
 	}
 
-	// In case we had a direct get subscription.
-	if mset.directSub != nil && stopping {
-		mset.unsubscribe(mset.directSub)
-		mset.directSub = nil
+	// In case we had a direct get subscriptions.
+	if stopping {
+		if mset.directSub != nil {
+			mset.unsubscribe(mset.directSub)
+			mset.directSub = nil
+		}
+		if mset.lastBySub != nil {
+			mset.unsubscribe(mset.lastBySub)
+			mset.lastBySub = nil
+		}
 	}
 
 	mset.active = false
@@ -3333,6 +3372,53 @@ func (mset *stream) processDirectGetRequest(_ *subscription, c *client, _ *Accou
 	}
 }
 
+// This is for direct get by last subject which is part of the subject itself.
+func (mset *stream) processDirectGetLastBySubjectRequest(_ *subscription, c *client, _ *Account, subject, reply string, rmsg []byte) {
+	_, msg := c.msgParts(rmsg)
+	if len(reply) == 0 {
+		return
+	}
+	// This version expects no payload.
+	if len(msg) != 0 {
+		hdr := []byte("NATS/1.0 408 Bad Request\r\n\r\n")
+		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
+		return
+	}
+	// Extract the key.
+	var key string
+	for i, n := 0, 0; i < len(subject); i++ {
+		if subject[i] == btsep {
+			if n == 4 {
+				if start := i + 1; start < len(subject) {
+					key = subject[i+1:]
+				}
+				break
+			}
+			n++
+		}
+	}
+	if len(key) == 0 {
+		hdr := []byte("NATS/1.0 408 Bad Request\r\n\r\n")
+		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
+		return
+	}
+
+	inlineOk := c.kind != ROUTER && c.kind != GATEWAY && c.kind != LEAF
+	if !inlineOk {
+		// Check how long we have been away from the readloop for the route or gateway or leafnode.
+		// If too long move to a separate go routine.
+		if elapsed := time.Since(c.in.start); elapsed < noBlockThresh {
+			inlineOk = true
+		}
+	}
+
+	if inlineOk {
+		mset.getDirectRequest(&JSApiMsgGetRequest{LastFor: key}, reply)
+	} else {
+		go mset.getDirectRequest(&JSApiMsgGetRequest{LastFor: key}, reply)
+	}
+}
+
 // Do actual work on a direct msg request.
 // This could be called in a Go routine if we are inline for a non-client connection.
 func (mset *stream) getDirectRequest(req *JSApiMsgGetRequest, reply string) {


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index da7da323..6e6022e1 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -18967,3 +18967,106 @@ func TestJetStreamMirrorUpdatesNotSupported(t *testing.T) {
 		t.Fatalf("Expected error %q, got %q", NewJSStreamMirrorNotUpdatableError(), err)
 	}
 }
+
+func TestJetStreamDirectGetBySubject(t *testing.T) {
+	conf := createConfFile(t, []byte(`
+		listen: 127.0.0.1:-1
+		jetstream: {max_mem_store: 64GB, max_file_store: 10TB}
+
+		ONLYME = {
+			publish = { allow = "$JS.API.DIRECT.GET.KV.vid.22.>"}
+		}
+
+		accounts: {
+			A: {
+				jetstream: enabled
+				users: [
+					{ user: admin, password: s3cr3t },
+					{ user: user, password: pwd, permissions: $ONLYME},
+				]
+			},
+		}
+	`))
+	defer removeFile(t, conf)
+
+	s, _ := RunServerWithConfig(conf)
+	if config := s.JetStreamConfig(); config != nil {
+		defer removeDir(t, config.StoreDir)
+	}
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s, nats.UserInfo("admin", "s3cr3t"))
+	defer nc.Close()
+
+	// Do by hand for now.
+	cfg := &StreamConfig{
+		Name:        "KV",
+		Storage:     MemoryStorage,
+		Subjects:    []string{"vid.*.>"},
+		MaxMsgsPer:  1,
+		AllowDirect: true,
+	}
+	addStream(t, nc, cfg)
+
+	// Add in mirror as well.
+	cfg = &StreamConfig{
+		Name:         "M",
+		Storage:      MemoryStorage,
+		Mirror:       &StreamSource{Name: "KV"},
+		MirrorDirect: true,
+	}
+	addStream(t, nc, cfg)
+
+	v22 := "vid.22.speed"
+	v33 := "vid.33.speed"
+	_, err := js.Publish(v22, []byte("100"))
+	require_NoError(t, err)
+	_, err = js.Publish(v33, []byte("55"))
+	require_NoError(t, err)
+
+	// User the restricted user.
+	nc, _ = jsClientConnect(t, s, nats.UserInfo("user", "pwd"))
+	defer nc.Close()
+
+	errCh := make(chan error, 10)
+	nc.SetErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, e error) {
+		select {
+		case errCh <- e:
+		default:
+		}
+	})
+
+	getSubj := fmt.Sprintf(JSDirectGetLastBySubjectT, "KV", v22)
+	m, err := nc.Request(getSubj, nil, time.Second)
+	require_NoError(t, err)
+	require_True(t, string(m.Data) == "100")
+
+	// Now attempt to access vid 33 data..
+	getSubj = fmt.Sprintf(JSDirectGetLastBySubjectT, "KV", v33)
+	_, err = nc.Request(getSubj, nil, 200*time.Millisecond)
+	require_Error(t, err) // timeout here.
+
+	select {
+	case e := <-errCh:
+		if !strings.HasPrefix(e.Error(), "nats: Permissions Violation") {
+			t.Fatalf("Expected a permissions violation but got %v", e)
+		}
+	case <-time.After(time.Second):
+		t.Fatalf("Expected to get a permissions error, got none")
+	}
+
+	// Now make sure mirrors are doing right thing with new way as well.
+	var sawMirror bool
+	getSubj = fmt.Sprintf(JSDirectGetLastBySubjectT, "KV", v22)
+	for i := 0; i < 100; i++ {
+		m, err := nc.Request(getSubj, nil, time.Second)
+		require_NoError(t, err)
+		if shdr := m.Header.Get(JSStream); shdr == "M" {
+			sawMirror = true
+			break
+		}
+	}
+	if !sawMirror {
+		t.Fatalf("Expected to see the mirror respond at least once")
+	}
+}