Skip to content

Commit 9856321

Browse files
xoofxCopilot
andcommitted
Enforce cumulative output size limits
Co-authored-by: Copilot <[email protected]>
1 parent 8180fb6 commit 9856321

File tree

5 files changed

+160
-12
lines changed

5 files changed

+160
-12
lines changed

src/Scriban.Tests/TestRuntime.cs

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,37 @@ public void String_And_Null_Concatenated_Should_Not_Null()
9090
Assert.AreEqual("my name is ", result);
9191
}
9292

93+
[Test]
94+
public void LimitToStringShouldApplyToCumulativeRenderOutput()
95+
{
96+
var context = new TemplateContext
97+
{
98+
LimitToString = 5
99+
};
100+
var template = Template.Parse("{{ 'abc' }}{{ 'def' }}");
101+
102+
var result = template.Render(context);
103+
104+
Assert.AreEqual("abcde...", result);
105+
}
106+
107+
[Test]
108+
public void ResetShouldClearCumulativeRenderOutputTracking()
109+
{
110+
var context = new TemplateContext
111+
{
112+
LimitToString = 5
113+
};
114+
var largeTemplate = Template.Parse("{{ 'abc' }}{{ 'def' }}");
115+
var smallTemplate = Template.Parse("{{ 'xy' }}");
116+
117+
Assert.AreEqual("abcde...", largeTemplate.Render(context));
118+
119+
context.Reset();
120+
121+
Assert.AreEqual("xy", smallTemplate.Render(context));
122+
}
123+
93124
[Test]
94125
public void TestAssignValToDictionary()
95126
{

src/Scriban/Functions/ObjectFunctions.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -511,7 +511,7 @@ static void WriteValue(TemplateContext context, Utf8JsonWriter writer, object va
511511
}
512512

513513
var type = value?.GetType() ?? typeof(object);
514-
var shouldTrackPath = value != null && !type.IsValueType && value is not string;
514+
var shouldTrackPath = value != null && !type.IsValueType && !(value is string);
515515
var addedToPath = false;
516516

517517
if (shouldTrackPath)

src/Scriban/ScribanAsync.generated.cs

Lines changed: 51 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -485,23 +485,35 @@ public async ValueTask<TemplateContext> WriteAsync(string text, int startIndex,
485485
// Write indents if necessary
486486
if (_previousTextWasNewLine)
487487
{
488-
await Output.WriteAsync(CurrentIndent, 0, CurrentIndent.Length, CancellationToken).ConfigureAwait(false);
488+
if (!await WriteOutputChunkAsync(CurrentIndent, 0, CurrentIndent.Length).ConfigureAwait(false))
489+
{
490+
return this;
491+
}
489492
_previousTextWasNewLine = false;
490493
}
491-
await Output.WriteAsync(text, index, indexEnd - index, CancellationToken).ConfigureAwait(false);
494+
if (!await WriteOutputChunkAsync(text, index, indexEnd - index).ConfigureAwait(false))
495+
{
496+
return this;
497+
}
492498
break;
493499
}
494500

495501
var length = newLineIndex - index;
496502
// Write indents if necessary
497503
if (_previousTextWasNewLine && (IndentOnEmptyLines || length != 0 && (length != 1 || text[index] != '\r')))
498504
{
499-
await Output.WriteAsync(CurrentIndent, 0, CurrentIndent.Length, CancellationToken).ConfigureAwait(false);
505+
if (!await WriteOutputChunkAsync(CurrentIndent, 0, CurrentIndent.Length).ConfigureAwait(false))
506+
{
507+
return this;
508+
}
500509
_previousTextWasNewLine = false;
501510
}
502511

503512
// We output the new line
504-
await Output.WriteAsync(text, index, length + 1, CancellationToken).ConfigureAwait(false);
513+
if (!await WriteOutputChunkAsync(text, index, length + 1).ConfigureAwait(false))
514+
{
515+
return this;
516+
}
505517
index = newLineIndex + 1;
506518
_previousTextWasNewLine = true;
507519
}
@@ -512,7 +524,7 @@ public async ValueTask<TemplateContext> WriteAsync(string text, int startIndex,
512524
{
513525
_previousTextWasNewLine = text[startIndex + count - 1] == '\n';
514526
}
515-
await Output.WriteAsync(text, startIndex, count, CancellationToken).ConfigureAwait(false);
527+
await WriteOutputChunkAsync(text, startIndex, count).ConfigureAwait(false);
516528
}
517529
}
518530

@@ -565,6 +577,40 @@ public async ValueTask<TemplateContext> WriteLineAsync()
565577
await WriteAsync(NewLine).ConfigureAwait(false);
566578
return this;
567579
}
580+
581+
private async ValueTask<bool> WriteOutputChunkAsync(string text, int startIndex, int count)
582+
{
583+
if (count <= 0)
584+
{
585+
return true;
586+
}
587+
588+
var allowedCount = GetAllowedOutputCount(count);
589+
if (allowedCount > 0)
590+
{
591+
await Output.WriteAsync(text, startIndex, allowedCount, CancellationToken).ConfigureAwait(false);
592+
_currentOutputLength += allowedCount;
593+
}
594+
595+
if (allowedCount < count)
596+
{
597+
await WriteOutputLimitEllipsisAsync().ConfigureAwait(false);
598+
return false;
599+
}
600+
601+
return true;
602+
}
603+
604+
private async ValueTask WriteOutputLimitEllipsisAsync()
605+
{
606+
if (_hasOutputLimitEllipsis)
607+
{
608+
return;
609+
}
610+
611+
await Output.WriteAsync("...", 0, 3, CancellationToken).ConfigureAwait(false);
612+
_hasOutputLimitEllipsis = true;
613+
}
568614
}
569615
}
570616

src/Scriban/TemplateContext.Helpers.cs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ public virtual IList ToList(SourceSpan span, object value)
7979

8080
private int _objectToStringLevel;
8181
private int _currentToStringLength;
82+
private int _currentOutputLength;
83+
private bool _hasOutputLimitEllipsis;
8284

8385
/// <summary>
8486
/// Called whenever an objects is converted to a string. This method can be overriden.
@@ -512,4 +514,4 @@ public virtual object ToObject(SourceSpan span, object value, Type destinationTy
512514
}
513515

514516
}
515-
}
517+
}

src/Scriban/TemplateContext.cs

Lines changed: 74 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,62 @@ public virtual TemplateContext Write(SourceSpan span, object textAsObject)
700700
return this;
701701
}
702702

703+
private void ResetOutputLimitTracking()
704+
{
705+
_currentOutputLength = 0;
706+
_hasOutputLimitEllipsis = false;
707+
}
708+
709+
private bool WriteOutputChunk(string text, int startIndex, int count)
710+
{
711+
if (count <= 0)
712+
{
713+
return true;
714+
}
715+
716+
var allowedCount = GetAllowedOutputCount(count);
717+
if (allowedCount > 0)
718+
{
719+
Output.Write(text, startIndex, allowedCount);
720+
_currentOutputLength += allowedCount;
721+
}
722+
723+
if (allowedCount < count)
724+
{
725+
WriteOutputLimitEllipsis();
726+
return false;
727+
}
728+
729+
return true;
730+
}
731+
732+
private int GetAllowedOutputCount(int requestedCount)
733+
{
734+
if (LimitToString <= 0)
735+
{
736+
return requestedCount;
737+
}
738+
739+
var remaining = LimitToString - _currentOutputLength;
740+
if (remaining <= 0)
741+
{
742+
return 0;
743+
}
744+
745+
return Math.Min(requestedCount, remaining);
746+
}
747+
748+
private void WriteOutputLimitEllipsis()
749+
{
750+
if (_hasOutputLimitEllipsis)
751+
{
752+
return;
753+
}
754+
755+
Output.Write("...", 0, 3);
756+
_hasOutputLimitEllipsis = true;
757+
}
758+
703759
/// <summary>
704760
/// Writes the text to the current <see cref="Output"/>
705761
/// </summary>
@@ -756,23 +812,35 @@ public TemplateContext Write(string text, int startIndex, int count)
756812
// Write indents if necessary
757813
if (_previousTextWasNewLine)
758814
{
759-
Output.Write(CurrentIndent, 0, CurrentIndent.Length);
815+
if (!WriteOutputChunk(CurrentIndent, 0, CurrentIndent.Length))
816+
{
817+
return this;
818+
}
760819
_previousTextWasNewLine = false;
761820
}
762-
Output.Write(text, index, indexEnd - index);
821+
if (!WriteOutputChunk(text, index, indexEnd - index))
822+
{
823+
return this;
824+
}
763825
break;
764826
}
765827

766828
var length = newLineIndex - index;
767829
// Write indents if necessary
768830
if (_previousTextWasNewLine && (IndentOnEmptyLines || length != 0 && (length != 1 || text[index] != '\r')))
769831
{
770-
Output.Write(CurrentIndent, 0, CurrentIndent.Length);
832+
if (!WriteOutputChunk(CurrentIndent, 0, CurrentIndent.Length))
833+
{
834+
return this;
835+
}
771836
_previousTextWasNewLine = false;
772837
}
773838

774839
// We output the new line
775-
Output.Write(text, index, length + 1);
840+
if (!WriteOutputChunk(text, index, length + 1))
841+
{
842+
return this;
843+
}
776844
index = newLineIndex + 1;
777845
_previousTextWasNewLine = true;
778846
}
@@ -782,7 +850,7 @@ public TemplateContext Write(string text, int startIndex, int count)
782850
if(count > 0){
783851
_previousTextWasNewLine = text[startIndex + count - 1] == '\n';
784852
}
785-
Output.Write(text, startIndex, count);
853+
WriteOutputChunk(text, startIndex, count);
786854
}
787855
}
788856

@@ -902,6 +970,7 @@ public virtual void Reset()
902970

903971
CachedTemplates.Clear();
904972
_memberAccessors.Clear();
973+
ResetOutputLimitTracking();
905974
}
906975

907976
/// <summary>

0 commit comments

Comments
 (0)